"""Where does the model know that its current answer is right? Linear read-out per layer.

At every probe cut of a --record-all shadow run (label: forced answer at that cut is correct), the residual stream of
every layer is recorded; for each layer an L2-regularised logistic regression is evaluated with trajectory-grouped
5-fold cross-validation (no trajectory in both train and test). Also tests a temporal feature: the change of the
state since the previous cut (consistency of the internal state over time).
usage: layer_probe.py --rec results/rec/raw_clean.jsonl --items "data/coc/*.jsonl" --exclude data/coc/mini_math.jsonl
"""
import argparse, glob, json, os, sys
import torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from transformers import AutoModelForCausalLM, AutoTokenizer
import transformers.integrations.sdpa_attention as _sdpa
_sdpa.use_gqa_in_sdpa = lambda *a, **k: False
from grade import grade_one

ap = argparse.ArgumentParser()
ap.add_argument("--rec", required=True)
ap.add_argument("--items", nargs="+", required=True)
ap.add_argument("--exclude", nargs="*", default=[])
ap.add_argument("--model", default="models/LFM3-1.7B")
ap.add_argument("--cache", default="results/halt/layer_feats.pt")
args = ap.parse_args()

if not os.path.exists(args.cache):
    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16, attn_implementation="sdpa",
                                                 trust_remote_code=True).cuda().eval()
    L = model.config.num_hidden_layers
    cap = {}
    for k, layer in enumerate(model.model.layers):
        layer.register_forward_hook(lambda m, i, o, k=k: cap.__setitem__(k, (o[0] if isinstance(o, tuple) else o)[0]))
    ends = {i for i in range(len(tok)) if (t := tok.convert_ids_to_tokens(i)) is not None
            and (t.endswith(".") or t.endswith("Ċ"))}
    items = {}
    for pat in args.items:
        for f in glob.glob(pat):
            for l in open(f, encoding="utf-8"):
                r = json.loads(l); items[r["id"]] = r
    excl = set()
    for f in args.exclude:
        excl |= {json.loads(l)["id"] for l in open(f, encoding="utf-8")}
    X, Y, T, TASK, IDX = [], [], [], [], []
    n = 0
    for l in open(args.rec, encoding="utf-8"):
        r = json.loads(l)
        if r["id"] in excl or r["id"] not in items or len(r.get("probe_log") or []) == 0:
            continue
        it = items[r["id"]]
        p_ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True,
                                        add_generation_prompt=True, return_dict=False)
        resp = tok(r["completion"][len("<think>"):], add_special_tokens=False)["input_ids"][:8192]
        pos, ys = [], []
        for p in r["probe_log"]:
            c = min(p["cut"], len(resp))
            while c > 1 and resp[c - 1] not in ends:
                c -= 1
            if c > 1:
                pos.append(len(p_ids) + c - 1)
                ys.append(float(grade_one({**r, "completion": "<think>x</think>" + p["text"], "n_tokens": 0}, it)["correct"]))
        if not pos:
            continue
        with torch.no_grad():
            model(input_ids=torch.tensor([p_ids + resp], device="cuda"), use_cache=False)
        X.append(torch.stack([cap[k][pos].float().cpu() for k in range(L)], 1))  # [n_pos, L, H]
        Y += ys; T += [n] * len(pos); TASK += [it["task"]] * len(pos); IDX += list(range(len(pos)))
        n += 1
    torch.save({"X": torch.cat(X).half(), "y": torch.tensor(Y), "traj": torch.tensor(T), "task": TASK,
                "idx": torch.tensor(IDX)}, args.cache)
d = torch.load(args.cache)
X, y, traj, idx = d["X"].float(), d["y"], d["traj"], d["idx"]
print(X.shape[0], "cut points,", int(y.sum()), "positive,", int(traj.max()) + 1, "trajectories")
prev = torch.zeros_like(X)
for i in range(1, len(X)):
    if traj[i] == traj[i - 1]:
        prev[i] = X[i - 1]
has_prev = idx > 0


def auc(s, yy):
    o = torch.argsort(s); r = torch.empty(len(s)); r[o] = torch.arange(1, len(s) + 1, dtype=torch.float)
    n1 = yy.sum(); n0 = len(yy) - n1
    return ((r[yy == 1].sum() - n1 * (n1 + 1) / 2) / (n1 * n0)).item()


def cv_auc(F, yy, tr_ids, lam):
    folds = torch.tensor([hash(int(t)) % 5 for t in tr_ids])
    scores = torch.zeros(len(yy))
    for f in range(5):
        te = folds == f; trn = ~te
        mu, sd = F[trn].mean(0), F[trn].std(0) + 1e-6
        A = (F - mu) / sd
        w = torch.zeros(A.shape[1], requires_grad=True); b = torch.zeros(1, requires_grad=True)
        opt = torch.optim.LBFGS([w, b], max_iter=200, line_search_fn="strong_wolfe")
        def closure():
            opt.zero_grad()
            loss = torch.nn.functional.binary_cross_entropy_with_logits(A[trn] @ w + b, yy[trn]) + lam * w.pow(2).sum()
            loss.backward(); return loss
        opt.step(closure)
        scores[te] = (A[te] @ w + b).detach()
    return auc(scores, yy)


print("layer  AUC(state)  AUC(state - previous cut)   [5-fold, grouped by trajectory]")
for k in range(0, X.shape[1], 2):
    a1 = max(cv_auc(X[:, k], y, traj, lam) for lam in (1e-1, 1.0))
    m = has_prev
    a2 = max(cv_auc(torch.cat([X[m, k], X[m, k] - prev[m, k]], 1), y[m], traj[m], lam) for lam in (1e-1, 1.0))
    print(f"{k:5d}  {a1:10.3f}  {a2:10.3f}", flush=True)