"""DiffuRefill-1B as an lm-eval-harness model, so it is scored on exactly the
tasks, prompts and metrics the autoregressive baselines are.

A masked-diffusion model has no left-to-right likelihood of its own; it scores a
continuation by the chain rule in left-to-right order. For a continuation
x_1..x_n after context c, row i is

    <bos> c x_1 .. x_{i-1} [MASK] * (n - i + 1)

and log p(x_i | c, x_<i) is read at the first masked position. The model knows
the continuation's length (the masks are there) but sees none of its later
tokens, so the sum over i is a proper likelihood of x under the decoding order
the model is later sampled in. The same chain gives `is_greedy`: every x_i is
the argmax at its own step.

Rows are grouped by exact length, so no padding is ever attended to (attention
is bidirectional), and only the scored position goes through the 73k-way head.
"""
import os, sys, time, collections
import torch
import torch.nn.functional as F

os.environ.setdefault("BG_DOCMASK", "0")
sys.path.insert(0, "/tmp")
from big_common import Net, MASK_ID, VOCAB, SEQ
from lm_eval.api.model import LM

REPO = "Asilarkness/DiffuRefill-1B"


class DLM(LM):
    def __init__(self, ckpt="final/DiffuRefill-1B-200k-swa.pt", budget=65536, device="cuda"):
        super().__init__()
        from huggingface_hub import hf_hub_download
        from transformers import AutoTokenizer
        self.tk = AutoTokenizer.from_pretrained("openbmb/MiniCPM4-0.5B", trust_remote_code=True)
        self.bos = self.tk.bos_token_id if self.tk.bos_token_id is not None else 1
        p = ckpt if ckpt.startswith("/") else hf_hub_download(REPO, ckpt, local_dir="/tmp/dlm")
        g = torch.load(p, map_location="cpu", weights_only=False)
        net = Net()
        net.load_state_dict({k: v.float() for k, v in g["model"].items()})
        self.net = net.to(device, torch.bfloat16).eval()
        self.dev, self.budget = device, budget

    def _enc(self, s):
        return [i for i in self.tk(s, add_special_tokens=False).input_ids if i < VOCAB - 2]

    def _pair(self, ctx, cont):
        # as lm-eval's HF model: trailing context whitespace belongs to the continuation
        n = len(ctx) - len(ctx.rstrip())
        if n:
            cont, ctx = ctx[-n:] + cont, ctx[:-n]
        whole = self._enc(ctx + cont)
        c = self._enc(ctx)
        x = whole[len(c):]
        if not x or whole[:len(c)] != c:            # tokenisation merged across the boundary
            x = self._enc(cont)
        c = [self.bos] + c
        c = c[-(SEQ - len(x)):]                     # keep the end of a long context
        return c, x

    @torch.no_grad()
    def _score_rows(self, rows):
        """rows: list of (ids, pos, target); all ids the same length. -> logp, argmax"""
        ids = torch.tensor([r[0] for r in rows], device=self.dev)
        pos = torch.tensor([r[1] for r in rows], device=self.dev)
        tgt = torch.tensor([r[2] for r in rows], device=self.dev)
        h = self.net.body(ids)[torch.arange(len(rows), device=self.dev), pos]
        lp = F.log_softmax(self.net.head(h).float()[:, :VOCAB - 2], -1)
        return lp.gather(1, tgt[:, None])[:, 0].tolist(), (lp.argmax(-1) == tgt).tolist()

    def loglikelihood(self, requests, disable_tqdm=False):
        if os.environ.get("EV_SCORER", "chain") == "elbo":
            return self._elbo(requests)
        jobs = collections.defaultdict(list)        # length -> rows (req, ids, pos, tgt)
        for k, r in enumerate(requests):
            c, x = self._pair(*r.args)
            n = len(x)
            for i in range(n):
                ids = c + x[:i] + [MASK_ID] * (n - i)
                jobs[len(ids)].append((k, ids, len(c) + i, x[i]))
        lp = [0.0] * len(requests)
        greedy = [True] * len(requests)
        t0, done, total = time.time(), 0, sum(len(v) for v in jobs.values())
        for L in sorted(jobs):
            rows = jobs[L]
            B = max(1, self.budget // L)
            for s in range(0, len(rows), B):
                part = rows[s:s + B]
                with torch.autocast("cuda", torch.bfloat16):
                    l, g = self._score_rows([(r[1], r[2], r[3]) for r in part])
                for r, a, b in zip(part, l, g):
                    lp[r[0]] += a
                    greedy[r[0]] &= b
                done += len(part)
            print("  ll rows %d/%d (%.0fs)" % (done, total, time.time() - t0), flush=True)
        return list(zip(lp, greedy))

    @torch.no_grad()
    def _elbo(self, requests, N=int(os.environ.get("EV_ELBO_N", "32"))):
        """LLaDA's conditional estimator: mask l of the n continuation tokens
        (l stratified over 1..n, positions random), score (n/l) * sum of the
        masked log-probs, average over N draws. A lower bound on log p(x|c);
        used only as a cross-check of the chain-rule numbers."""
        import random
        rng = random.Random(0)
        lp = [0.0] * len(requests)
        jobs = collections.defaultdict(list)        # length -> (req, ids, mask positions, targets, weight)
        for k, r in enumerate(requests):
            c, x = self._pair(*r.args)
            n = len(x)
            for j in range(N):
                l = 1 + min(n - 1, int((j + rng.random()) / N * n))
                pos = rng.sample(range(n), l)
                ids = c + [MASK_ID if i in pos else t for i, t in enumerate(x)]
                jobs[len(ids)].append((k, ids, [len(c) + i for i in pos], [x[i] for i in pos], n / l / N))
        for L in sorted(jobs):
            rows = jobs[L]
            B = max(1, self.budget // L)
            for s in range(0, len(rows), B):
                part = rows[s:s + B]
                ids = torch.tensor([r[1] for r in part], device=self.dev)
                with torch.autocast("cuda", torch.bfloat16):
                    h = self.net.body(ids)
                for b, r in enumerate(part):
                    q = F.log_softmax(self.net.head(h[b, r[2]]).float()[:, :VOCAB - 2], -1)
                    lp[r[0]] += r[4] * float(q.gather(1, torch.tensor(r[3], device=self.dev)[:, None]).sum())
        return [(v, False) for v in lp]

    def loglikelihood_rolling(self, requests, disable_tqdm=False):
        raise NotImplementedError("perplexity tasks are not used")

    def generate_until(self, requests, disable_tqdm=False):
        raise NotImplementedError("generation tasks are not used")


if __name__ == "__main__":
    import json, lm_eval
    tasks = os.environ.get("EV_TASKS", "hellaswag,arc_easy,arc_challenge,piqa,winogrande,openbookqa,sciq,lambada_openai,mmlu").split(",")
    out = os.environ.get("EV_OUT", "/tmp/bench_dlm.json")
    m = DLM(os.environ.get("EV_CKPT", "final/DiffuRefill-1B-200k-swa.pt"))
    res = {}
    for t in tasks:                                  # one at a time: a broken task must not sink the rest
        try:
            lim = int(os.environ["EV_LIMIT"]) if os.environ.get("EV_LIMIT") else None
            r = lm_eval.simple_evaluate(model=m, tasks=[t], num_fewshot=0, bootstrap_iters=0, limit=lim)
            res.update({k: v for k, v in r["results"].items()})
            print("RESULT", t, json.dumps(r["results"].get(t, {})), flush=True)
        except Exception as e:
            print("TASKFAIL", t, str(e)[:200], flush=True)
        json.dump(res, open(out, "w"))
    print("BENCH DONE", flush=True)