"""Compare the four per-loop-target mean Jacobians in loop-relative coordinates.

For loop k (target t_k = 48k-1) and physical layer L, take J_k[L] = J_{v->t_k}
with v = 48(k-1)+L. Matched-L cross-loop pairs share distance-to-target
(47 - L), so this comparison is distance-controlled by construction.
"""

import json

import torch
import torch.nn.functional as F

from jlens.lens import JacobianLens

torch.set_grad_enabled(False)
dev = "cuda" if torch.cuda.is_available() else "cpu"

LENSES = {
    1: "jacobian_lens_ouro-2.6b_t47.pt",
    2: "jacobian_lens_ouro-2.6b_t95.pt",
    3: "jacobian_lens_ouro-2.6b_t143.pt",
    4: "jacobian_lens_ouro-2.6b.pt",
}
lens = {k: JacobianLens.load(p) for k, p in LENSES.items()}
d = lens[4].jacobians[0].shape[0]
LS = list(range(47))
M = {k: {L: lens[k].jacobians[48 * (k - 1) + L].float().to(dev) for L in LS}
     for k in lens}
print({k: str(lens[k]) for k in lens}, flush=True)


def vcos(a, b):
    return F.cosine_similarity(a.flatten(), b.flatten(), 0).item()


R = 32
svd = {}
shape = {}
for k in M:
    for L in LS:
        J = M[k][L]
        U, S, Vh = torch.linalg.svd(J)
        svd[k, L] = (U[:, :R].contiguous(), S.cpu(), Vh[:R].contiguous())
        p = S / S.sum()
        fro = J.norm().item()
        shape[k, L] = {
            "fro": fro / d ** 0.5,
            "smax": S[0].item(),
            "erank": torch.exp(-(p * p.log()).sum()).item(),
            "stable_rank": (fro ** 2 / S[0].item() ** 2),
            "id_cos": (torch.trace(J) / (d ** 0.5 * fro)).item(),
            "sym_frac": (((J + J.T) / 2).norm() ** 2 / fro ** 2).item(),
            "diag_frac": (J.diagonal().norm() ** 2 / fro ** 2).item(),
        }
print("svd done", flush=True)

pairs = [(a, b) for a in range(1, 5) for b in range(a + 1, 5)]
res = {"matched": {}, "sub_r": {}, "sub_l": {}, "pairwise": {}}
for L in LS:
    cs, sr, sl = [], [], []
    for a, b in pairs:
        c = vcos(M[a][L], M[b][L])
        cs.append(c)
        res["pairwise"].setdefault(f"{a}-{b}", {})[L] = c
        Va, Vb = svd[a, L][2], svd[b, L][2]
        Ua, Ub = svd[a, L][0], svd[b, L][0]
        sr.append(((Va @ Vb.T).norm() ** 2 / R).item())
        sl.append(((Ua.T @ Ub).norm() ** 2 / R).item())
    res["matched"][L] = sum(cs) / len(cs)
    res["sub_r"][L] = sum(sr) / len(sr)
    res["sub_l"][L] = sum(sl) / len(sl)

adj = [vcos(M[k][L], M[k][L + 1]) for k in M for L in LS[:-1]]
off = {}
for dL in (2, 4, 8):
    xs = [vcos(M[a][L], M[b][L + dL]) for a, b in pairs for L in LS
          if L + dL in LS]
    off[dL] = sum(xs) / len(xs)
g = torch.Generator().manual_seed(0)
rnd = []
for _ in range(300):
    a, b = pairs[torch.randint(6, (1,), generator=g).item()]
    L1 = torch.randint(47, (1,), generator=g).item()
    L2 = torch.randint(47, (1,), generator=g).item()
    if L1 != L2:
        rnd.append(vcos(M[a][L1], M[b][L2]))

print("\n=== cross-loop matched-L (distance-controlled) ===", flush=True)
print("  L | mcos  subR  subL |  erank k=1..4        | idcos k=1..4", flush=True)
for L in LS:
    er = "  ".join(f"{shape[k, L]['erank']:5.0f}" for k in range(1, 5))
    ic = "  ".join(f"{shape[k, L]['id_cos']:.2f}" for k in range(1, 5))
    print(f" {L:3d} | {res['matched'][L]:.2f}  {res['sub_r'][L]:.2f}  "
          f"{res['sub_l'][L]:.2f} | {er} | {ic}", flush=True)

band = [L for L in LS if 24 <= L <= 40]
outb = [L for L in LS if L < 24]
print("\n=== summary ===", flush=True)
print(f"matched-L cos: all {sum(res['matched'].values())/47:.3f}  "
      f"band(24-40) {sum(res['matched'][L] for L in band)/len(band):.3f}  "
      f"early(<24) {sum(res['matched'][L] for L in outb)/len(outb):.3f}", flush=True)
print(f"baselines: within-loop adjacent-L {sum(adj)/len(adj):.3f}  "
      f"cross-loop offset dL=2 {off[2]:.3f} dL=4 {off[4]:.3f} dL=8 {off[8]:.3f}  "
      f"random {sum(rnd)/len(rnd):.3f}", flush=True)
print(f"subspace overlap r={R} (random baseline {R/d:.3f}): "
      f"right {sum(res['sub_r'].values())/47:.3f}  "
      f"left {sum(res['sub_l'].values())/47:.3f}", flush=True)
print("per-pair mean matched-L cos:", flush=True)
for a, b in pairs:
    m = sum(res["pairwise"][f"{a}-{b}"].values()) / 47
    print(f"  loop{a}-loop{b}: {m:.3f}", flush=True)

print("\n=== chain-rule factorisation: J(v->191) vs J(mid->191) @ J(v->mid) ===",
      flush=True)
res["compose"] = {}
for mid, km in ((95, 2), (143, 3)):
    Jm = lens[4].jacobians[mid].float().to(dev)
    for v in (8, 24, 32, 40, 56, 72, 80, 88, 104, 120, 128, 136):
        if v >= mid:
            continue
        Jv = lens[km].jacobians[v].float().to(dev)
        direct = lens[4].jacobians[v].float().to(dev)
        comp = Jm @ Jv
        c = vcos(direct, comp)
        nr = (comp.norm() / direct.norm()).item()
        res["compose"][f"{v}->{mid}->191"] = {"cos": c, "norm_ratio": nr}
        print(f"  v{v:3d} ->{mid}->191: cos {c:.3f}  "
              f"|comp|/|direct| {nr:.2f}", flush=True)

json.dump(
    {"matched": res["matched"], "sub_r": res["sub_r"], "sub_l": res["sub_l"],
     "pairwise": res["pairwise"], "compose": res["compose"],
     "baselines": {"adjacent": sum(adj) / len(adj),
                   "offset": off, "random": sum(rnd) / len(rnd)},
     "shape": {f"{k},{L}": shape[k, L] for k, L in shape},
     "spectra_top16": {f"{k},{L}": svd[k, L][1][:16].tolist()
                       for k, L in svd}},
    open("jcompare_analysis.json", "w"), default=float)
print("\nDONE_JCOMPARE", flush=True)