"""Gap analysis: transport quality vs distance on Ouro, using mid-band targets
(v120, v132) to place proven-content sources (loop-2 band v72-88) at distances
32-60, plus the loop-end lenses (t143, t191) for the rest of the curve.

Metrics (matched to qwen_overlay.py): ||J||, state-space cosine J h_v vs h_t,
KL and top-1 agreement against the target layer's own readout. Plus the
semantic test: does the transported loop-2 band state carry the Italy
direction at the mid-band target?
"""

import json

import torch
import torch.nn.functional as F

from jlens.hooks import ActivationRecorder
from jlens.lens import JacobianLens
from ouro_lens import from_ouro, load_ouro

torch.set_grad_enabled(False)

hf, tok = load_ouro()
hf = hf.cuda()
model = from_ouro(hf, tok)
NP, V = model.n_physical_layers, model.n_layers
d = model.d_model

LENSES = {
    120: "jacobian_lens_ouro_mid_t120.pt",
    132: "jacobian_lens_ouro_mid_t132.pt",
    143: "jacobian_lens_ouro-2.6b_t143.pt",
    191: "jacobian_lens_ouro-2.6b.pt",
}
lenses = {t: JacobianLens.load(p) for t, p in LENSES.items()}
J = {t: {v: lenses[t].jacobians[v].cuda() for v in lenses[t].source_layers}
     for t in LENSES}
for t in LENSES:
    print(t, lenses[t], flush=True)

with open("prompts_1000.json") as f:
    eval_prompts = json.load(f)[980:1000]

results = {"norms": {}, "curves": {}}
for t in LENSES:
    results["norms"][t] = {v: (lenses[t].jacobians[v].norm().item() / d**0.5)
                           for v in lenses[t].source_layers}

acc = {t: {"cos": {}, "kl": {}, "agree": {}} for t in LENSES}
n_pos = 0
for pi, p in enumerate(eval_prompts):
    ids = model.encode(p, max_length=128)
    with ActivationRecorder(model.layers, at=list(range(V))) as rec:
        model.forward(ids)
    acts = {v: rec.activations[v].detach()[0] for v in range(V)}
    seq = ids.shape[1]
    positions = list(range(16, seq - 1, 8))
    n_pos += len(positions)
    for t in LENSES:
        h_t = acts[t][positions].float()
        t_logits = model.unembed(h_t).float()
        t_lsm = t_logits.log_softmax(-1)
        t_top = t_logits.argmax(-1)
        for v in J[t]:
            h = acts[v][positions].float()
            th = h @ J[t][v].T
            c = F.cosine_similarity(th, h_t, dim=-1).sum().item()
            lg = model.unembed(th).float()
            k = F.kl_div(lg.log_softmax(-1), t_lsm, log_target=True,
                         reduction="sum").item()
            a = (lg.argmax(-1) == t_top).float().sum().item()
            acc[t]["cos"][v] = acc[t]["cos"].get(v, 0) + c
            acc[t]["kl"][v] = acc[t]["kl"].get(v, 0) + k
            acc[t]["agree"][v] = acc[t]["agree"].get(v, 0) + a
    print(f"prompt {pi + 1}/20", flush=True)

for t in LENSES:
    results["curves"][t] = {
        m: {v: acc[t][m][v] / n_pos for v in acc[t][m]} for m in acc[t]
    }

print("\n=== state-space cosine (J h_v vs h_t) by distance ===", flush=True)
print("dist |  t120  t132  t143  t191", flush=True)
for dist in (1, 4, 8, 16, 24, 32, 40, 44, 48, 56, 60, 72):
    row = []
    for t in LENSES:
        v = t - dist
        row.append(f"{results['curves'][t]['cos'].get(v, float('nan')):5.2f}"
                   if v in results["curves"][t]["cos"] else "  -  ")
    print(f" {dist:3d} | " + "  ".join(row), flush=True)

print("\n=== loop-2 band (v72-88, proven content) through each target ===", flush=True)
for t in (120, 132, 143, 191):
    cs = [results["curves"][t]["cos"][v] for v in range(72, 89)
          if v in results["curves"][t]["cos"]]
    dists = f"{t-88}-{t-72}"
    print(f"  ->t{t}: dist {dists:7s} mean cos {sum(cs)/len(cs):.3f}", flush=True)

# ---- semantic test: Italy through the mid-band transport ----
prompt = "Fact: The currency used in the country shaped like a boot is"
ids = model.encode(prompt)
with ActivationRecorder(model.layers, at=list(range(V))) as rec:
    model.forward(ids)
acts = {v: rec.activations[v].detach()[0] for v in range(V)}
pos = -2
U = hf.lm_head.weight.detach().float().cuda()
tok_italy = tok(" Italy", add_special_tokens=False).input_ids[0]
tok_ctrl = tok(" France", add_special_tokens=False).input_ids[0]

print("\n=== Italy-direction projection at mid-band targets (boot position) ===", flush=True)
for t in (120, 132):
    d_it = J[143][t].T @ U[tok_italy]
    d_it = d_it / d_it.norm()
    d_fr = J[143][t].T @ U[tok_ctrl]
    d_fr = d_fr / d_fr.norm()
    h_t = acts[t][pos].float()
    ref_it, ref_fr = float(h_t @ d_it), float(h_t @ d_fr)
    print(f"  target v{t}: actual state  proj_Italy={ref_it:7.1f}  proj_France={ref_fr:7.1f}", flush=True)
    for v in (72, 76, 80, 84, 88):
        th = J[t][v] @ acts[v][pos].float()
        print(f"    from v{v} (dist {t-v:2d}): proj_Italy={float(th @ d_it):7.1f}  "
              f"proj_France={float(th @ d_fr):7.1f}  "
              f"|Jh|/|h_t|={float(th.norm()/h_t.norm()):.2f}", flush=True)
    top = model.unembed(J[143][t] @ (J[t][80] @ acts[80][pos].float())).topk(5).indices
    print(f"    composed readout v80->v{t}->t143: "
          f"{[tok.decode([x]).strip() for x in top.tolist()]}", flush=True)

json.dump({str(t): {m: {str(v): results["curves"][t][m][v] for v in results["curves"][t][m]}
                    for m in results["curves"][t]} for t in results["curves"]}
          | {"norms": {str(t): {str(v): n for v, n in results["norms"][t].items()}
                       for t in results["norms"]}},
          open("gap_analysis.json", "w"))
print("\nDONE_GAP_ANALYSIS", flush=True)