"""Cross-env gradient fingerprint: three gradient paths on the SAME ckpt (s2000) + SAME seeded batch.
Run under torch 2.10 (local A6000) and torch 2.3.1 (timan107 GTX1080); diverging fingerprints pinpoint
which path is numerically broken on the pascal env (suspects: double-backward reg paths).
  (a) plain CE backward through the feedforward block   (bp-style backward)
  (b) resreg path: tforce graph at z_T1 + autograd.grad (single backward through force graph)
  (c) jacreg path: autograd.functional.jvp(create_graph) + backward   (double-backward)
Prints total grad-norm per path + first-3 per-tensor norms."""
import torch, torch.nn.functional as F
import lt_ep_train as L

torch.manual_seed(0)
blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick'); blk.qknorm = True
# RANDOM-INIT probe: no ckpt load (the sick regime is step 0-400)
torch.manual_seed(42)
idx, y = L.get_batch('train', 8, 256)
xin = blk.embed(idx).detach()
zs = L.relax(blk, xin.clone(), xin, 150, 0.1)  # init operator: res ~e-9 regime

def report(tag, params, grads):
    tot = sum(float(g.pow(2).sum()) for g in grads if g is not None) ** 0.5
    firsts = " ".join(f"{float(g.norm()):.3e}" if g is not None else "None" for g in grads[:3])
    fin = all(torch.isfinite(g).all() for g in grads if g is not None)
    print(f"{tag}: total={tot:.4e} finite={fin} first3=[{firsts}]", flush=True)

# (a) plain BP backward
x = blk.embed(idx)
h1 = F.layer_norm(x, (blk.C,), blk.ln1g, blk.ln1b)
h2 = F.layer_norm(x, (blk.C,), blk.ln2g, blk.ln2b)
h = x + blk.attn(h1) + (F.gelu(h2 @ blk.fc + blk.fcb, approximate='tanh') @ blk.pj + blk.pjb)
loss = F.cross_entropy((h @ blk.Wh).reshape(-1, L.vocab), y.reshape(-1))
ga = torch.autograd.grad(loss, blk.allp, allow_unused=True)
report("(a) plain-BP  ", blk.allp, list(ga))

# (b) resreg path
with torch.enable_grad():
    Fz = blk.tforce(zs, xin)
    Rr = (0.1 * Fz).pow(2).sum() / (zs.pow(2).sum() + 1e-9)
    gb = torch.autograd.grad(Rr, blk.block, allow_unused=True)
report("(b) resreg    ", blk.block, list(gb))

# (c) jacreg double-backward path
er = torch.randn_like(zs)
with torch.enable_grad():
    Jv = torch.autograd.functional.jvp(blk.nc_force, zs.detach(), er, create_graph=True)[1]
    R = 0.1 * (Jv ** 2).sum() / (er ** 2).sum()
    gc = torch.autograd.grad(R, blk.block, allow_unused=True)
report("(c) jacreg dbb", blk.block, list(gc))
# (d) the FULL ep_step gradient (holo track estimator + AEP + regs) — the actual training signal
blk.track = True
ge, res = L.ep_step(blk, idx, y, 150, 20, 0.1, 0.02, jacreg=0.1, holo=2, hr=0.02,
                    t1max=300, res_est=1e-4, t2sel=40, corr_every=1, res_gate=0.0, resreg=0.2)
gs = [g for g in ge.values() if g is not None]
tot = sum(float(g.pow(2).sum()) for g in gs) ** 0.5
fin = all(torch.isfinite(g).all() for g in gs)
print(f"(d) full ep_step: total={tot:.4e} finite={fin} res={res:.2e} n={len(gs)}", flush=True)
print("GRAD_ENV_PROBE_DONE", flush=True)