"""Router leakage / generalization check.

The router was trained on LongBench examples [0:N] (disjoint train/eval slices). Here we load the SAVED
router and evaluate its gold-block coverage on a FRESH set of LongBench examples the training never saw
(offset far past N). If coverage on unseen docs ~ training-time eval coverage, there is no train/eval
leakage or overfit. Also checks the serving condition: features come from the resident summary attention,
which IS available at sparse serving (no train-only signal).
"""
from __future__ import annotations
import argparse, glob, json, torch, torch.nn as nn
import scripts.e2e_nll as N
import scripts.train_router_real as TR
LB = "/work/hdd/bdjx/hxia3/hf_cache/hub/datasets--Syon-Li--LongbenchSeg/snapshots/*/longbench_segmented.jsonl"


def load_offset(tok, sink, max_len, skip, take):
    """Load LongBench QA, SKIPPING the first `skip` valid examples (training set), take next `take`."""
    out, seen = [], 0
    for line in open(glob.glob(LB)[0]):
        if len(out) >= take:
            break
        if len(line) > max_len * 40:
            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
        seen += 1
        if seen <= skip:               # skip the examples used during router training
            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:9]):
            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) >= 3 and ce0 < len(ids)-1:
            out.append((ids, blk, snk, kl0, ce0, "lb"))
    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("--skip", type=int, default=400)
    ap.add_argument("--take", type=int, default=60); ap.add_argument("--max-len", type=int, default=1200)
    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_offset(tok, sink, args.max_len, args.skip, args.take)
    print(f"FRESH examples (skipped first {args.skip} training docs): {len(ex)}")
    data = TR.collect(model, ex, summ, dev, args.max_len)     # (feat, oracle_top1) via dense forward
    ks = [1, 2, 3]
    cov = {k: 0.0 for k in ks}; base = {k: 0.0 for k in ks}; tot = 0
    for feat, tgt in data:
        sc = router(feat.to(dev)).squeeze(-1).cpu()
        rank = sc.argsort(-1, descending=True)
        mrank = feat.mean(-1).argsort(-1, descending=True)     # untrained baseline
        for k in ks:
            cov[k] += (rank[:, :k] == tgt.view(-1, 1)).any(-1).float().sum().item()
            base[k] += (mrank[:, :k] == tgt.view(-1, 1)).any(-1).float().sum().item()
        tot += feat.shape[0]
    print(f"\ncoverage of the model's oracle top-1 block on UNSEEN docs ({tot} answer positions):")
    print(f"{'k':>3} | {'untrained mean-summary':>22} | {'trained router':>14}")
    for k in ks:
        print(f"{k:>3} | {base[k]/tot:>22.3f} | {cov[k]/tot:>14.3f}")
    print("\nIf trained >> untrained AND ~matches training-time eval (0.96@k=2), the router generalizes")
    print("to unseen documents -> no train/eval leakage, no overfit.")


if __name__ == "__main__":
    main()