"""Adjoint-truncation hypothesis: the cos(EP,BPTT) ceiling at near-edge operators is the finite-T2
window (adjoint needs ~1/|Re mu| ~ 50-500 steps there; deep-contraction ops converge fast -> 0.98).
Sweep the tracking window t2sel in {40, 80, 160, 320} at s2000, everything else fixed (bsub=4, T1=150,
3 batches). Prediction: cos climbs with window; the climb rate quantifies the truncation bias that the
2.40-plateau memory called the 'estimator bias-floor'."""
import torch
import lt_ep_train as L
from diag_cos import cos_ep_bptt

torch.manual_seed(0)
blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick'); blk.qknorm = True
ck = torch.load('runs/redx_traj/s2000.pt', map_location=L.dev)
with torch.no_grad():
    for p, w in zip(blk.allp, ck['allp']):
        p.copy_(w.to(L.dev))
blk.track = True
torch.manual_seed(11)
batches = [L.get_batch('train', 24, 256) for _ in range(3)]

for w in (40, 80, 160, 320):
    cs = []
    for idx, y in batches:
        c, _ = cos_ep_bptt(blk, idx, y, 150, 20, 0.1, 0.02, holo=2, hr=0.02, t2sel=w, bsub=4)
        cs.append(c)
    print(f"t2sel={w:<4} cos={' '.join(f'{c:.4f}' for c in cs)}  mean={sum(cs)/len(cs):.4f}", flush=True)
print("T2_PROBE_DONE", flush=True)