"""Honest accuracy-retention test: per-example regression analysis, not aggregate F1.

For each real LongBench QA: is dense correct? is sparse correct? Then count REGRESSIONS (dense right ->
sparse wrong) vs GAINS (dense wrong -> sparse right). Aggregate F1 can hide regressions; this can't.
Correct = gold answer (normalized) contained in the generated first line.
"""
from __future__ import annotations
import argparse, glob, json, string, torch, torch.nn as nn
import scripts.demo_infer as DM
LB = "/work/hdd/bdjx/hxia3/hf_cache/hub/datasets--Syon-Li--LongbenchSeg/snapshots/*/longbench_segmented.jsonl"


def norm(s):
    return " ".join("".join(c for c in s.lower() if c not in string.punctuation).split())


def load(tok, sink, max_len, n, max_ans=50):
    out = []
    for line in open(glob.glob(LB)[0]):
        if len(out) >= n: 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"]) < 3: continue
        ans = r["answers"][0]
        if not (1 <= len(ans) <= max_ans): 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][:600], -1, sp=True)
        for bi,c in enumerate(r["chunks"][1:11]):
            for s in sink: ids.append(s); blk.append(bi); snk.append(True)
            add("\n"+c, bi)
        q0=len(ids); add(f"\nQuestion: {r['input']}\nAnswer:", -2)
        if len(ids) <= max_len and (max(blk)+1) >= 3:
            out.append((ids, blk, snk, q0, ans, r.get("dataset","?"), r["input"], list(r["chunks"][1:11])))
    return out


@torch.no_grad()
def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--model", required=True)
    ap.add_argument("--router", required=True); ap.add_argument("--n", type=int, default=80)
    ap.add_argument("--k", type=int, default=2); ap.add_argument("--max-len", type=int, default=2000)
    ap.add_argument("--max-new", type=int, default=24); ap.add_argument("--dump", default=None)
    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(tok, sink, args.max_len, args.n)
    print(f"per-example regression test on {len(ex)} short-answer LongBench QA (k={args.k}) ...\n")

    dh=sh=reg=gain=both=0; regs=[]; dump=[]
    for ids, blk, snk, q0, gold, ds, question, docs in ex:
        gd = DM.gen(model,tok,ids,blk,snk,q0,'dense',None,0,args.max_new,0,summ).split("Answer:")[-1].split("\n")[0]
        gs = DM.gen(model,tok,ids,blk,snk,q0,'sparse',router,args.k,args.max_new,0,summ).split("Answer:")[-1].split("\n")[0]
        dc = norm(gold) in norm(gd); sc = norm(gold) in norm(gs)
        dh+=dc; sh+=sc; both+=(dc and sc)
        if dc and not sc: reg+=1; regs.append((gold, gd.strip()[:50], gs.strip()[:50], ds))
        if sc and not dc: gain+=1
        dump.append({"dataset": ds, "question": question, "documents": docs, "gold": gold,
                     "dense_output": gd.strip(), "sparse_output": gs.strip(),
                     "dense_correct": bool(dc), "sparse_correct": bool(sc)})
    n=len(ex)
    if args.dump:
        import json as J
        J.dump({"model": args.model, "k": args.k, "summary_tokens": summ,
                "summary": {"n": n, "dense_correct": dh, "sparse_correct": sh,
                            "regressions": reg, "gains": gain}, "examples": dump},
               open(args.dump, "w"), indent=2, ensure_ascii=False)
        print(f"dumped {n} examples -> {args.dump}")
    print(f"dense correct : {dh}/{n} ({dh/n*100:.0f}%)")
    print(f"sparse correct: {sh}/{n} ({sh/n*100:.0f}%)")
    print(f"REGRESSIONS (dense right -> sparse wrong): {reg}/{n}")
    print(f"GAINS       (dense wrong -> sparse right): {gain}/{n}")
    print(f"net change: {sh-dh:+d}   (agreement both-correct: {both})")
    if regs:
        print("\nregression cases (gold | dense | sparse | task):")
        for g,d,s,ds in regs[:8]:
            print(f"  gold={g!r} | dense={d!r} | sparse={s!r} | {ds}")


if __name__ == "__main__":
    main()