"""Train the landmark router on REAL RAG data (LongbenchSeg), distilling to the model's own attention.

Fixes the synthetic-capitals distribution gap. No gold-block labels needed: the target at each answer
position is the ORACLE top-1 block = argmax of full-block attention mass there (the model's own
retrieval signal). Router = linear over per-(layer,head) landmark attention -> predict that block.

Run: source env.sh && HF_HOME=/projects/bdjx/hxia3/hf_cache_proj PYTHONPATH=. \
     python scripts/train_router_real.py --model hf_release/model --save-router checkpoints/qwen3-14b-router-real.pt
"""
from __future__ import annotations

import argparse
import random

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

import scripts.distill_lora as D    # load_data (real LongbenchSeg/SemanticSeg), masks
import scripts.e2e_nll as N         # landmarks_content, mask_for


@torch.no_grad()
def collect(model, examples, summary_n, dev, max_tokens=1100):
    """Per example -> (feat [A,nb,L*H], oracle_top1 [A]). feat = per-(layer,head) attention to each
    block's SUMMARY = its last `summary_n` content tokens (mean). summary_n=1 -> old single-landmark."""
    data = []
    for ids, blk, sink, kl0, ce0, src in examples:
        if len(ids) > max_tokens:
            continue
        nb = max(blk) + 1
        if nb < 3:
            continue
        n = len(ids)
        rows = {i: set(range(nb)) for i in range(n - kl0)}
        mask = N.mask_for(blk, sink, kl0, nb, rows, dev)
        out = model(input_ids=torch.tensor([ids], device=dev), attention_mask=mask,
                    output_attentions=True)
        atts = torch.stack(out.attentions, 0)[:, 0].float()    # [L,H,n,n]
        _, isc = N.landmarks_content(blk, sink, nb)
        bpos = {b: [] for b in range(nb)}
        content = torch.zeros(nb, n, device=dev)
        for j, c in enumerate(isc):
            if c:
                content[blk[j], j] = 1.0; bpos[blk[j]].append(j)
        a0 = ce0
        if a0 >= n - 1:
            continue
        A = n - a0
        # feat[t, b] = mean over block b's last-summary_n content tokens of attn, per (layer,head)
        feat = torch.zeros(A, nb, atts.shape[0] * atts.shape[1])
        for b in range(nb):
            cols = torch.tensor(bpos[b][-summary_n:], device=dev)
            v = atts[:, :, a0:n, :][:, :, :, cols].mean(-1)    # [L,H,A]
            feat[:, b, :] = v.permute(2, 0, 1).reshape(A, -1).cpu()
        full_mass = (atts.mean(dim=(0, 1))[a0:n] @ content.T)
        data.append((feat, full_mass.argmax(-1).cpu()))
    return data


def coverage(scorer, data, ks, nb_cap):
    cov = {k: 0.0 for k in ks}; tot = 0
    for feat, tgt in data:
        sc = scorer(feat)
        rank = sc.argsort(-1, descending=True)
        for k in ks:
            cov[k] += (rank[:, :k] == tgt.view(-1, 1)).any(-1).float().sum().item()
        tot += feat.shape[0]
    return {k: cov[k] / tot for k in ks}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--n-examples", type=int, default=350)
    ap.add_argument("--max-len", type=int, default=1100)
    ap.add_argument("--epochs", type=int, default=400)
    ap.add_argument("--wd", type=float, default=1e-3)
    ap.add_argument("--summary-tokens", type=int, default=8, help="block summary = last N content tokens")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--save-router", default=None)
    args = ap.parse_args()
    dev = "cuda"

    print(f"loading {args.model} ...")
    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(
        args.model, dtype=torch.bfloat16, attn_implementation="eager").to(dev).eval()
    sink_ids = tok("\n", add_special_tokens=False)["input_ids"] * 4

    print("loading real RAG data (LongbenchSeg) ...")
    ex = D.load_data(tok, sink_ids, args.max_len, args.n_examples, args.seed)
    ex = [e for e in ex if e[5] == "lb"]                       # QA (query+answer) examples
    random.Random(1).shuffle(ex)
    print(f"  {len(ex)} QA examples")
    print(f"collecting features (summary = last {args.summary_tokens} tokens) + oracle targets ...")
    data = collect(model, ex, args.summary_tokens, dev, args.max_len)
    print(f"  usable: {len(data)} examples, {sum(f.shape[0] for f,_ in data)} answer positions")
    n_eval = max(20, len(data) // 5)
    ev, tr = data[:n_eval], data[n_eval:]
    LH = tr[0][0].shape[-1]; nb_max = max(f.shape[1] for f, _ in data)
    print(f"  feat dim L*H={LH} | train {len(tr)} / eval {len(ev)}")

    router = nn.Linear(LH, 1).to(dev)
    opt = torch.optim.Adam(router.parameters(), lr=1e-3, weight_decay=args.wd)
    Xtr = [(f.to(dev), t.to(dev)) for f, t in tr]
    for epq in range(1, args.epochs + 1):
        opt.zero_grad(); tl = 0.0
        for feat, tgt in Xtr:
            sc = router(feat).squeeze(-1)
            loss = F.cross_entropy(sc, tgt); loss.backward(); tl += loss.item()
        opt.step()
        if epq % 100 == 0 or epq == 1:
            print(f"epoch {epq:4d} | CE {tl/len(Xtr):.4f}")

    ks = [1, 2, 3]
    with torch.no_grad():
        rt = lambda f: router(f.to(dev)).squeeze(-1).cpu()
        mean_lm = lambda f: f.mean(-1)
        ev_r = coverage(rt, ev, ks, nb_max)
        ev_m = coverage(mean_lm, ev, ks, nb_max)
    print(f"\nheld-out coverage of the model's own oracle top-1 block (real LongBench QA):")
    print(f"{'k':>3} | {'untrained mean-lm':>17} | {'trained router (eval)':>21}")
    for k in ks:
        print(f"{k:>3} | {ev_m[k]:>17.3f} | {ev_r[k]:>21.3f}")

    if args.save_router:
        torch.save({"state_dict": router.state_dict(), "in_dim": LH,
                    "summary_tokens": args.summary_tokens,
                    "note": f"linear over per-(layer,head) attn to block's last-{args.summary_tokens} tokens; real LongBench QA"},
                   args.save_router)
        print(f"saved router -> {args.save_router}  (summary_tokens={args.summary_tokens})")


if __name__ == "__main__":
    main()