"""Box head as a speculative drafter for coordinate tokens (Jacobi verification over a shared prefix).

The box head answers every question of an image in one branch pass over the cached image prefix. Its boxes are
rendered as the backbone's own coordinate-token answer ("x0, y0, x1, y1]" in 0–1000 units, after a teacher-forced
JSON preamble) and verified in parallel by the backbone's LM head, with the adapter disabled (base model):

  pass k:  argmax g_t of every draft position given the draft before it, all questions batched as branches.
           Accept the longest prefix with g_t == draft_t; the token after it, g_n, is exact too.
           Next draft = accepted + g_n + g_{n+1..} (Jacobi iterate). Stop when nothing changes.

The fixed point is exactly the greedy decode under the same computation. Rows still unconverged after max_passes
finish with plain greedy decoding from their accepted prefix, so the output is always the greedy answer; the cost
is the number of sequential passes. max_passes=0 is plain greedy decoding (the coordinate-token control on the same
prompt, cache and kernels).

accept < 1 relaxes the check (typical acceptance, as in Medusa): a draft token also passes when the verifier gives it
at least accept x the probability of its argmax. The output is then no longer the greedy answer but the head's box
unless the verifier clearly disagrees; accept=1 is the exact rule above.
"""
from __future__ import annotations

import re

import numpy as np
import torch

from locus.models.boxvlm import BOX_TOKEN, BoxVLM

# Continues BoxVLM.prefix ("...Locate:"): the zero-shot grounding prompt of the Qwen harness, with a colon.
VERIFY_PROMPT = ' "{}" in the image and output its bbox coordinates in JSON format.'
PREAMBLE = '```json\n[\n\t{"bbox_2d": ['


class DraftVerify:
    def __init__(self, model: BoxVLM, max_passes: int = 4, max_tokens: int = 24, verifier_adapter: bool = False,
                 accept: float = 1.0):
        assert model.n_layers_run == model.n_layers_total, "the verifier needs the full LM (no early exit)"
        self.m, self.tok = model, model.ad.tok
        self.max_passes, self.max_tokens = max_passes, max_tokens
        self.close = self.tok.encode("]", add_special_tokens=False)[0]
        self.vsuffix = model.suffix.replace(BOX_TOKEN, "") + PREAMBLE
        self.verifier_adapter = verifier_adapter  # aLoRA trained with --ar-weight: adapter on verifier branches too
        self.log_accept = float(np.log(accept)) if accept < 1 else None
        self._closes = {}

    # -- token helpers ------------------------------------------------------
    def is_close(self, t: int) -> bool:
        if t not in self._closes:
            self._closes[t] = self.tok.decode([t]).lstrip().startswith("]")
        return self._closes[t]

    def cut(self, ids: list[int]) -> list[int]:
        """Truncate after the first token that closes the coordinate list; append "]" if there is none."""
        for i, t in enumerate(ids):
            if self.is_close(t):
                return ids[: i + 1]
        return ids + [self.close]

    def render(self, box) -> list[int]:
        return self.tok.encode(", ".join(str(int(round(min(max(float(v), 0.0), 1.0) * 1000))) for v in box) + "]",
                               add_special_tokens=False)

    def parse(self, ids: list[int]) -> np.ndarray:
        m = re.match(r"\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\]", self.tok.decode(ids))
        if m is None:
            return np.full(4, np.nan)
        x0, y0, x1, y1 = (min(int(v), 1000) / 1000 for v in m.groups())
        return np.array([min(x0, x1), min(y0, y1), max(x0, x1), max(y0, y1)])

    def verifier_inputs(self, texts, boxes=None):
        """Verifier-branch prompt tokens per question (continuing BoxVLM.prefix), and rendered answers if boxes given."""
        prompts = [self.tok.encode(VERIFY_PROMPT.format(t) + self.vsuffix, add_special_tokens=False) for t in texts]
        return prompts, None if boxes is None else [self.render(b) for b in boxes]

    # -- verifier passes ----------------------------------------------------
    def _greedy_at(self, state, prompts, drafts, rows, row_image):
        """One batched branch pass: argmax at every draft position (and the one after the last) for the given rows.
        With accept < 1, a draft token the verifier finds plausible enough is returned in place of the argmax."""
        seqs = [prompts[i] + drafts[i] for i in rows]
        L = max(map(len, seqs))
        ids = torch.full((len(seqs), L), self.tok.pad_token_id, dtype=torch.long)
        mask = torch.zeros(len(seqs), L, dtype=torch.long)
        for j, s in enumerate(seqs):
            ids[j, : len(s)], mask[j, : len(s)] = torch.tensor(s), 1
        dev = row_image.device
        if self.m.fast is not None:
            h = self.m.branch_pass(state, ids.to(dev), mask.to(dev), row_image[rows], keep_cache=True,
                                   adapter=self.verifier_adapter).last_hidden_state
        else:
            ctx = self.m.ad.model.disable_adapter() if self.m.lora_r and not self.verifier_adapter else torch.no_grad()
            with ctx:
                h = self.m.branch_pass(state, ids.to(dev), mask.to(dev), row_image[rows], keep_cache=True).last_hidden_state
        spans = [(j, len(prompts[i]) - 1, len(prompts[i]) + len(drafts[i])) for j, i in enumerate(rows)]
        sel = torch.cat([h[j, a:b] for j, a, b in spans])  # LM head only where needed: len(draft) + 1 per row
        logits = self.m.base.lm_head(sel).float()
        pred = logits.argmax(-1)
        if self.log_accept is not None:
            draft = torch.tensor([t for i in rows for t in drafts[i] + [-1]], device=pred.device)
            lp = logits.log_softmax(-1)
            ok = (draft >= 0) & (lp.gather(1, draft.clamp(min=0)[:, None])[:, 0] - lp.max(-1).values >= self.log_accept)
            pred = torch.where(ok, draft, pred)
        pred = pred.tolist()
        out, k = [], 0
        for _, a, b in spans:
            out.append(pred[k : k + b - a])
            k += b - a
        return out

    @torch.no_grad()
    def __call__(self, prefix_inputs, texts: list[str], row_image: torch.Tensor, verify: bool = True,
                 draft: np.ndarray | None = None) -> dict:
        """`draft` (boxes from another model's head) replaces this model's head as the drafter."""
        m = self.m
        state = m.prefill(prefix_inputs)
        bid, bmask = m.tokenize_branches(texts)
        dev = row_image.device
        head, conf = m.forward_shared(None, bid.to(dev), bmask.to(dev), row_image, state=state)
        head = head.float().cpu().numpy()
        res = dict(head=head, conf=torch.sigmoid(conf.float()).cpu().numpy())
        if not verify:
            return res
        assert m.alora or not m.lora_r or self.verifier_adapter, \
            "the base-model verifier needs a base prefix: use an aLoRA (or frozen) head, or a jointly trained adapter"
        prompts = self.verifier_inputs(texts)[0]
        R = len(texts)
        drafts = [self.render(b) for b in (head if draft is None else draft)] if self.max_passes else [[] for _ in range(R)]
        acc = [0] * R  # length of the greedy-exact prefix of each draft
        passes = np.zeros(R, int)
        active = list(range(R)) if self.max_passes else []
        for _ in range(self.max_passes):
            if not active:
                break
            g = self._greedy_at(state, prompts, drafts, active, row_image)
            nxt = []
            for i, gi in zip(active, g):
                passes[i] += 1
                d = drafts[i]
                n = next((k for k, (a, b) in enumerate(zip(d, gi)) if a != b and not (self.is_close(a) and self.is_close(b))), len(d))
                if n == len(d):  # fixed point: every draft token is the greedy token
                    acc[i] = len(d)
                    continue
                acc[i] = n + 1  # g_n is exact as well
                drafts[i] = self.cut(d[:n] + gi[n:])
                nxt.append(i)
            active = nxt
        # Unconverged rows (and max_passes=0): plain greedy decoding from the exact prefix.
        todo = [i for i in range(R) if acc[i] < len(drafts[i]) or not drafts[i]]
        for i in todo:
            drafts[i] = drafts[i][: acc[i]]
        while todo:
            g = self._greedy_at(state, prompts, [d for d in drafts], todo, row_image)
            nxt = []
            for i, gi in zip(todo, g):
                passes[i] += 1
                drafts[i] = drafts[i] + [gi[-1]]
                if not self.is_close(gi[-1]) and len(drafts[i]) < self.max_tokens:
                    nxt.append(i)
            todo = nxt
        res.update(box=np.stack([self.parse(d) for d in drafts]), passes=passes, tokens=np.array([len(d) for d in drafts]))
        return res