"""Gate for --fastfp (Anderson z*) and --bf16polish: endpoint parity + residual + eval-count/time vs the
150-step Euler reference, on two operators (deeply-trained fast-adaptive + near-edge s2000 — the hard
case: Anderson historically fails on cycling ops, s2000's marginal band is the stress test).
GPU shared with abl_pair — timings are relative, parity/evals exact."""
import time, torch
import lt_ep_train as L

for name, path in (('s2000', 'runs/redx_traj/s2000.pt'), ('fast-adaptive', 'runs/ep_fast_adaptive.pt')):
    torch.manual_seed(0)
    blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick'); blk.qknorm = True
    ck = torch.load(path, map_location=L.dev)
    with torch.no_grad():
        for p, w in zip(blk.allp, ck['allp']):
            p.copy_(w.to(L.dev))
    torch.manual_seed(42)
    idx, _ = L.get_batch('train', 24, 256)
    xin = blk.embed(idx).detach()

    torch.cuda.synchronize(); t = time.time()
    z_ref = L.relax(blk, xin.clone(), xin, 150, 0.1)
    torch.cuda.synchronize(); t_ref = time.time() - t
    r_ref = (L.relax(blk, z_ref, xin, 1, 0.1) - z_ref).norm().item()
    print(f"[{name}] euler150 : {t_ref:5.2f}s  res={r_ref:.2e}", flush=True)

    torch.cuda.synchronize(); t = time.time()
    z_aa, evals = L.anderson_relax(blk, xin.clone(), xin, 150, 0.1)
    torch.cuda.synchronize(); t_aa = time.time() - t
    r_aa = (L.relax(blk, z_aa, xin, 1, 0.1) - z_aa).norm().item()
    zd = ((z_aa - z_ref).norm() / (z_ref.norm() + 1e-12)).item()
    print(f"[{name}] anderson : {t_aa:5.2f}s  res={r_aa:.2e}  evals={evals}  z-diff={zd:.2e}", flush=True)

    blk.bf16polish = 20
    torch.cuda.synchronize(); t = time.time()
    z_bf = L.relax(blk, xin.clone(), xin, 150, 0.1)
    torch.cuda.synchronize(); t_bf = time.time() - t
    blk.bf16polish = 0
    r_bf = (L.relax(blk, z_bf, xin, 1, 0.1) - z_bf).norm().item()
    zdb = ((z_bf - z_ref).norm() / (z_ref.norm() + 1e-12)).item()
    print(f"[{name}] bf16+20  : {t_bf:5.2f}s  res={r_bf:.2e}  z-diff={zdb:.2e}", flush=True)
print("FASTFP_GATE_DONE", flush=True)