"""Set Block Decoding inference wrapper for Qwen2 (Belumind).

Adds sbd_generate(): lossless speculative decoding driven by three proposal
sources -- the causal NTP frontier, MATP over a trailing masked block, and
n-gram lookup over the prompt -- with exact verification folded into the next
proposal forward pass.

Reference implementation. It runs without a KV cache: every step recomputes the
whole sequence. tokens/forward is an accounting metric, not a wall-clock claim.
See the model card.
"""

import torch
from transformers.models.qwen2.modeling_qwen2 import Qwen2ForCausalLM

DEFAULT_MASK_TOKEN_ID = 151662  # <|fim_pad|> in the Qwen2.5 vocabulary


def sliding_block_mask(total_len, block_start, device, dtype):
    """Causal over the prefix, bidirectional from block_start onward.

    The block slides with the frontier instead of sitting on a fixed grid, so
    the last committed position always has a purely causal context and a valid
    NTP prediction always exists.
    """
    idx = torch.arange(total_len, device=device)
    causal = idx.view(1, total_len) <= idx.view(total_len, 1)
    inblk = idx >= block_start
    allowed = causal | (inblk.view(1, total_len) & inblk.view(total_len, 1))
    m = torch.zeros(total_len, total_len, device=device, dtype=dtype)
    m.masked_fill_(~allowed, torch.finfo(dtype).min)
    return m.view(1, 1, total_len, total_len)


def ngram_propose(ids, k, n_max=4, window=3000):
    """Most recent match of the longest suffix, and whatever followed it."""
    n_ids = len(ids)
    for n in range(n_max, 1, -1):
        if n_ids < n:
            continue
        pat = ids[-n:]
        lo = max(0, n_ids - window)
        for i in range(n_ids - n - 1, lo - 1, -1):
            if ids[i:i + n] == pat:
                cand = ids[i + n:i + n + k]
                if cand:
                    return cand
    return []


class SBDQwen2ForCausalLM(Qwen2ForCausalLM):
    """Qwen2 with a Set Block Decoding entrypoint. generate() is unaffected."""

    @torch.no_grad()
    def sbd_generate(self, input_ids, attention_mask=None, max_new_tokens=128,
                     block_size=8, threshold=0.6,
                     sources=("ntp", "matp", "ngram"),
                     mask_token_id=None, eos_token_id=None,
                     stop_at_eos=True, return_stats=False, **kwargs):
        if input_ids.dim() != 2 or input_ids.shape[0] != 1:
            raise ValueError("sbd_generate supports batch size 1 only")
        if attention_mask is not None and int(attention_mask.sum()) != attention_mask.numel():
            raise ValueError("padded prompts are not supported")

        dev = input_ids.device
        dt = next(self.parameters()).dtype
        if mask_token_id is None:
            mask_token_id = getattr(self.config, "sbd_mask_token_id", DEFAULT_MASK_TOKEN_ID)
        if eos_token_id is None:
            eos_token_id = self.config.eos_token_id
        if not stop_at_eos or eos_token_id is None:
            eos_ids = set()
        elif isinstance(eos_token_id, int):
            eos_ids = {eos_token_id}
        else:
            eos_ids = set(eos_token_id)

        use_matp = "matp" in sources
        use_ngram = "ngram" in sources
        B = max(1, int(block_size))

        p0 = input_ids.shape[1]
        seq = input_ids.clone()
        verified = p0
        fwd = 0
        rollbacks = 0
        acc = {"ntp": 0, "matp": 0, "ngram": 0}
        pend = []
        stop_at = None
        cap = int(max_new_tokens) + 4

        while fwd < cap:
            bs = seq.shape[1]
            block = torch.full((1, B), mask_token_id, device=dev, dtype=seq.dtype)
            inp = torch.cat([seq, block], 1)
            mask = sliding_block_mask(inp.shape[1], bs, dev, dt)
            logits = self(input_ids=inp, attention_mask=mask, use_cache=False).logits[0]
            fwd += 1

            # ---- verification. Free: this forward computed these conditionals
            # anyway in order to make the next proposal.
            if bs > verified:
                pred = logits[verified - 1:bs - 1].argmax(-1)
                bad = (pred != seq[0, verified:bs]).nonzero()
                if len(bad):
                    k = int(bad[0])
                    rollbacks += 1
                    for j in range(max(0, k - 1)):
                        acc[pend[j]] += 1
                    acc["ntp"] += 1 if k == 0 else 2
                    seq = torch.cat([seq[:, :verified + k], pred[k].view(1, 1)], 1)
                    verified = verified + k + 1
                    pend = []
                    stop_at = self._first_eos(seq, p0, verified, eos_ids)
                    if stop_at is not None or verified - p0 >= max_new_tokens:
                        break
                    # the rest of this forward described a prefix that was just
                    # rejected; it cannot be used to propose. Take a fresh one.
                    continue
                acc["ntp"] += 1
                for j in range(len(pend)):
                    acc[pend[j]] += 1
                verified = bs
                pend = []
                stop_at = self._first_eos(seq, p0, verified, eos_ids)
                if stop_at is not None:
                    break

            if verified - p0 >= max_new_tokens:
                break

            # ---- proposal
            ntp = int(logits[bs - 1].argmax())
            nxt = [ntp]
            src = []
            if use_matp or use_ngram:
                conf, arg = logits[bs:bs + B].float().softmax(-1).max(-1)
                matp = []
                if use_matp:
                    j = 1
                    while j < B and float(conf[j]) >= threshold:
                        matp.append(int(arg[j]))
                        j += 1
                if use_ngram:
                    ngm = ngram_propose(seq[0].tolist() + [ntp], B - 1)
                    if len(matp) < B - 1 and len(ngm) > len(matp):
                        extra = ngm[len(matp):B - 1]
                    else:
                        extra = []
                else:
                    extra = []
                nxt += matp
                src += ["matp"] * len(matp)
                nxt += extra
                src += ["ngram"] * len(extra)
            pend = src
            seq = torch.cat([seq, torch.tensor([nxt], device=dev, dtype=seq.dtype)], 1)

        committed = verified - p0
        end = verified if stop_at is None else stop_at
        end = min(end, p0 + max_new_tokens)
        out = seq[:, :end]
        if not return_stats:
            return out
        tot = sum(acc.values())
        stats = {
            "forwards": fwd,
            "tokens_committed": committed,
            "tokens_returned": end - p0,
            "tokens_per_forward": (committed / fwd) if fwd else 0.0,
            "rollbacks": rollbacks,
            "accepted": dict(acc),
            "pct_from_ngram": (acc["ngram"] / tot) if tot else 0.0,
            "pct_from_matp": (acc["matp"] / tot) if tot else 0.0,
        }
        return out, stats

    @staticmethod
    def _first_eos(seq, p0, verified, eos_ids):
        if not eos_ids:
            return None
        for i in range(p0, verified):
            if int(seq[0, i]) in eos_ids:
                return i + 1
        return None