"""Is the name still IN the field at the answer position?  Frozen recall-bench checkpoint, eval mode,
stream of recall episodes (held-out filler half); at the first answer byte collect α (32x32 field) and
fit a ridge one-vs-rest linear probe (name, 16-way) on 1200 episodes, test on 600, per distance bucket.
Chance = 6.25 %.  Separates 'memory retains the information' from 'readout learned to use it'.
usage: probe_recall_alpha.py OUT.json CKPT [CKPT ...]"""
import json, random, sys
from pathlib import Path
import numpy as np
import torch
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import longmem_bench as LB  # noqa: E402
sys.path.insert(0, str(LB.ROOT / "tools"))
from chat_atom_native import load_model  # noqa: E402


@torch.no_grad()
def collect(m, eps):
    m.eval(); m.reset_state(reset_atomizer=True); X, y, d = [], [], []; tick = 0
    for text, ans, dist, name in eps:
        pk = m.atomizer.encode_bytes(text, reset=False)
        for i in range(len(pk)):
            m.forward_packet(pk[i]); tick += 1
            if tick % 64 == 0: LB.MB.flush(m)
            if i + 1 == ans:
                X.append(m.core.state.alpha.detach().reshape(-1).numpy().copy()); y.append(LB.NAMES.index(name)); d.append(dist)
    return np.array(X), np.array(y), np.array(d)


def ridge_probe(Xtr, ytr, Xte, lam=1.0):
    mu, sd = Xtr.mean(0), Xtr.std(0) + 1e-6
    A, B = (Xtr - mu) / sd, (Xte - mu) / sd
    Y = np.eye(16)[ytr]
    W = np.linalg.solve(A.T @ A + lam * len(A) * np.eye(A.shape[1]), A.T @ Y)
    return (B @ W).argmax(1)


def main():
    torch.set_num_threads(1)
    out, cks = sys.argv[1], sys.argv[2:]
    src = LB.filler_source(); ev = src[len(src) // 2:]
    rng = random.Random(4242)
    eps = [LB.episode(rng, ev) for _ in range(1800)]
    res = {}
    for ck in cks:
        if ck.startswith("fresh"):     # untrained controls: "fresh" (λ .05) / "fresh_slow" (+ slow block .25)
            import argparse
            a = argparse.Namespace(seed=20260913, leak=0.05, slow_block=0.25 if ck == "fresh_slow" else 0.0,
                                   slow_cap=0.02, slow_gain="ema", k=0, ladrop=0.5)
            m = LB.build(a)
        else:
            m = load_model(ck)
        X, y, d = collect(m, eps)
        pred = ridge_probe(X[:1200], y[:1200], X[1200:], lam=1e-2)
        yt, dt = y[1200:], d[1200:]
        r = {"overall_acc": float((pred == yt).mean())}
        for lo, hi, b in ((0, 32, "d<32"), (32, 64, "d32-63"), (64, 128, "d64-127"), (128, 999, "d128+")):
            sel = (dt >= lo) & (dt < hi)
            r[b] = {"n": int(sel.sum()), "acc": float((pred[sel] == yt[sel]).mean()) if sel.any() else None}
        res[ck] = r; print(ck, json.dumps(r), flush=True)
    Path(out).write_text(json.dumps(res, indent=1))


if __name__ == "__main__":
    main()