"""Block-diffusion inference for this checkpoint.
Usage:
    from blockdiff_inference import load, generate
    model, tok, mask_id = load("REPO_ID")
    print(generate(model, tok, mask_id, "The capital of France is"))
"""
import sys, torch, torch.nn as nn, torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

BLOCK_LEN = 32
CENTERED_CONV = True
SHIFT = True
MASK_TOKEN = "<|mdm_mask|>"


def _mask(attention_mask, B, T, device, window=None):
    i = torch.arange(T, device=device)
    m = (i[None, :] // BLOCK_LEN) <= (i[:, None] // BLOCK_LEN)
    if window is not None:
        m &= (i[:, None] - i[None, :]).abs() < window
    m = m[None, None].expand(B, 1, T, T).clone()
    if attention_mask is not None and attention_mask.dim() == 2:
        m &= attention_mask[:, -T:].bool()[:, None, None, :]
    return m | torch.eye(T, dtype=torch.bool, device=device)[None, None]


def _mask_fn(sliding):
    def fn(*args, **kw):
        emb = kw.get("input_embeds", kw.get("inputs_embeds"))
        if emb is None:
            emb = next(a for a in args if torch.is_tensor(a) and a.dim() == 3)
        win = (getattr(kw.get("config"), "sliding_window", None) or 256) if sliding else None
        return _mask(kw.get("attention_mask"), emb.shape[0], emb.shape[1], emb.device, win)
    return fn


class CenteredConv1d(nn.Module):
    def __init__(self, conv):
        super().__init__(); self.conv = conv
        L = conv.kernel_size[0]; self.lp, self.rp = (L - 1) // 2, L - 1 - (L - 1) // 2
        conv.padding = (0,)
    @property
    def weight(self): return self.conv.weight
    @property
    def bias(self): return self.conv.bias
    def forward(self, x): return self.conv(F.pad(x, (self.lp, self.rp)))


def load(repo_id, device="cuda", dtype=torch.bfloat16):
    tok = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(repo_id, dtype=dtype, trust_remote_code=True,
                                                 attn_implementation="sdpa").to(device).eval()
    for t in {sys.modules[type(m).__module__] for m in model.modules()}:
        for name in dir(t):
            lo = name.lower()
            if "mask" in lo and "causal" in lo and callable(getattr(t, name)):
                setattr(t, name, _mask_fn("sliding" in lo))
    for m in model.modules():
        if hasattr(m, "is_causal"): m.is_causal = False
    if CENTERED_CONV:
        for parent in list(model.modules()):
            for cn, ch in list(parent.named_children()):
                if isinstance(ch, nn.Conv1d) and ch.groups == ch.in_channels and ch.kernel_size[0] > 1:
                    setattr(parent, cn, CenteredConv1d(ch))
    model.config.use_cache = False
    return model, tok, tok.convert_tokens_to_ids(MASK_TOKEN)


def _logits(model, x):
    lg = model(input_ids=x).logits
    return torch.cat([lg[:, :1], lg[:, :-1]], 1) if SHIFT else lg


@torch.inference_mode()
def generate(model, tok, mask_id, prompt, gen_len=256, tau=0.9, temperature=0.0, repetition_penalty=1.1):
    x = tok(prompt, return_tensors="pt").input_ids.to(model.device)
    P, L, eos = x.shape[1], BLOCK_LEN, tok.eos_token_id
    produced, n = 0, (L - P % L if P % L else L)
    while produced < gen_len:
        x = torch.cat([x, torch.full((1, n), mask_id, device=x.device)], 1)
        s, e = x.shape[1] - n, x.shape[1]
        for _ in range(L):
            is_mask = x[0, s:e] == mask_id
            if not is_mask.any(): break
            lg = _logits(model, x)[0, s:e].float(); lg[:, mask_id] = -float("inf")
            if repetition_penalty != 1.0 and s > P:
                seen = x[0, P:s].unique(); v = lg[:, seen]
                lg[:, seen] = torch.where(v > 0, v / repetition_penalty, v * repetition_penalty)
            pr = F.softmax(lg / temperature if temperature > 0 else lg, -1)
            pred = torch.multinomial(pr, 1)[:, 0] if temperature > 0 else pr.argmax(-1)
            conf = pr.gather(-1, pred[:, None])[:, 0].masked_fill(~is_mask, -1.0)
            sel = (conf >= tau) & is_mask
            if not sel.any(): sel = torch.zeros_like(is_mask); sel[conf.argmax()] = True
            x[0, s:e] = torch.where(sel, pred, x[0, s:e])
        produced += n; n = L
        if (x[0, s:e] == eos).any(): break
    out = x[0, P:]
    hit = (out == eos).nonzero()
    if len(hit): out = out[: hit[0, 0]]
    return tok.decode(out, skip_special_tokens=True)