"""Real-scale benchmark: 50 LongBench RAG samples, real block sizes. Measures (1) quality recovery
(gold-answer NLL: dense vs static vs last-8 router-sparse) and (2) active-context savings (fraction of
block-content tokens attended per decode step). Uses the trained last-8 router for per-position select.
"""
from __future__ import annotations
import argparse, glob, json, torch, torch.nn as nn
import scripts.e2e_nll as N

LB = "/work/hdd/bdjx/hxia3/hf_cache/hub/datasets--Syon-Li--LongbenchSeg/snapshots/*/longbench_segmented.jsonl"


def load_lb(tok, sink, max_len, max_blocks, n_want):
    """Real LongBench QA with up to max_blocks real chunks (real block sizes), <= max_len tokens."""
    out = []
    for line in open(glob.glob(LB)[0]):
        if len(out) >= n_want:
            break
        if len(line) > max_len * 60:
            continue
        try:
            r = json.loads(line)
        except Exception:
            continue
        if not r.get("chunks") or not r.get("answers") or len(r["chunks"]) < 4:
            continue
        ids, blk, snk = [], [], []
        def add(t, b, sp=False, sk=False):
            for x in ([t] if sk else tok(t, add_special_tokens=sp)["input_ids"]):
                ids.append(x); blk.append(b); snk.append(sk)
        add(r["chunks"][0][:400], -1, sp=True)
        for bi, c in enumerate(r["chunks"][1:1+max_blocks]):
            for s in sink:
                ids.append(s); blk.append(bi); snk.append(True)
            add("\n" + c, bi)
        kl0 = len(ids); add(f"\nQuestion: {r['input']}\nAnswer:", -2)
        ce0 = len(ids); add(" " + r["answers"][0], -2)
        if len(ids) <= max_len and (max(blk)+1) >= 4 and ce0 < len(ids)-1:
            out.append((ids, blk, snk, kl0, ce0, "lb"))
    return out


def last8_feat(atts, bp, a0, n, nb, summ, dev):
    """[A, nb, L*H] = per-(layer,head) attn from each answer pos to each block's last-`summ` tokens."""
    A = n - a0
    feat = torch.zeros(A, nb, atts.shape[0]*atts.shape[1])
    for b in range(nb):
        cols = torch.tensor(bp[b][-summ:], device=dev)
        feat[:, b, :] = atts[:, :, a0:n, :][:, :, :, cols].mean(-1).permute(2,0,1).reshape(A,-1).cpu()
    return feat


@torch.no_grad()
def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--model", required=True)
    ap.add_argument("--router", required=True); ap.add_argument("--k", type=int, default=2)
    ap.add_argument("--n", type=int, default=50); ap.add_argument("--max-len", type=int, default=3500)
    ap.add_argument("--max-blocks", type=int, default=16)
    args = ap.parse_args(); dev="cuda"
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16,
              attn_implementation="eager", device_map="cuda").eval()
    ck = torch.load(args.router, map_location=dev); summ = ck["summary_tokens"]
    router = nn.Linear(ck["in_dim"],1).to(dev); router.load_state_dict(ck["state_dict"]); router.eval()
    sink = tok("\n", add_special_tokens=False)["input_ids"]*4

    ex = load_lb(tok, sink, args.max_len, args.max_blocks, args.n)
    print(f"benchmarking {len(ex)} real LongBench QA samples (k={args.k}, summary={summ}, <= {args.max_blocks} blocks) ...")
    nll = {m: [] for m in ["dense","static","router"]}
    fracs=[]; nblocks=[]; blk_tok=[]; ctx_tok=[]
    for ids, blk, snk, kl0, ce0, src in ex:
        n=len(ids); nb=max(blk)+1
        if nb<3 or ce0>=n-1: continue
        bp = D_block_pos(blk, snk, nb)
        clen = {b: len(bp[b]) for b in range(nb)}; total=sum(clen.values())
        dense_rows={i:set(range(nb)) for i in range(n-kl0)}
        out = model(input_ids=torch.tensor([ids],device=dev),
                    attention_mask=N.mask_for(blk,snk,kl0,nb,dense_rows,dev), output_attentions=True)
        atts = torch.stack(out.attentions,0)[:,0].float()
        A = n-ce0
        feat = last8_feat(atts, bp, ce0, n, nb, summ, dev)
        sc = router(feat.to(dev)).squeeze(-1).cpu()                     # [A, nb]
        rsel = {t: set(sc[t].argsort(descending=True)[:args.k].tolist()) for t in range(A)}
        s0 = set(sc[0].argsort(descending=True)[:args.k].tolist())
        def rows(sel_fn):
            r={}
            for row in range(kl0, n):
                r[row-kl0] = set(range(nb)) if row<ce0 else sel_fn(row-ce0)
            return r
        nll["dense"].append(N.gold_nll(model, ids, N.mask_for(blk,snk,kl0,nb,{i:set(range(nb)) for i in range(n-kl0)},dev), ce0, dev).mean().item())
        nll["static"].append(N.gold_nll(model, ids, N.mask_for(blk,snk,kl0,nb,rows(lambda t:s0),dev), ce0, dev).mean().item())
        nll["router"].append(N.gold_nll(model, ids, N.mask_for(blk,snk,kl0,nb,rows(lambda t:rsel[t]),dev), ce0, dev).mean().item())
        # savings: attended content fraction per answer pos = (active content + nb*summ)/total
        fr = [(sum(clen[b] for b in rsel[t]) + sum(min(summ,clen[b]) for b in range(nb)))/max(total,1) for t in range(A)]
        fracs.append(sum(fr)/len(fr)); nblocks.append(nb); blk_tok.append(total/nb); ctx_tok.append(total)
    mean=lambda x: sum(x)/len(x)
    d,s,r = mean(nll["dense"]), mean(nll["static"]), mean(nll["router"])
    print(f"\n=== quality recovery (gold-answer NLL over {len(fracs)} samples) ===")
    print(f"  dense (all blocks)      : {d:.3f}")
    print(f"  static (freeze top-{args.k})   : {s:.3f}")
    print(f"  router last-{summ} (top-{args.k})   : {r:.3f}   -> recovers {(s-r)/(s-d)*100:.0f}% of the static->dense gap")
    print(f"\n=== active-context savings ===")
    print(f"  avg blocks/sample: {mean(nblocks):.1f} | avg block size: {mean(blk_tok):.0f} tok | avg context: {mean(ctx_tok):.0f} tok")
    print(f"  avg block-content tokens attended per decode step: {mean(fracs)*100:.0f}%  -> ~{1/mean(fracs):.1f}x less context read")


def D_block_pos(blk, snk, nb):
    bp={b:[] for b in range(nb)}
    for j,(b,s) in enumerate(zip(blk,snk)):
        if b>=0 and not s: bp[b].append(j)
    return bp


if __name__=="__main__":
    main()