import json

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import torch

res = json.load(open("ouro_analysis.json"))
ll = json.load(open("ouro_probes_ll.json"))
NP, R = 48, 4
V = NP * R
bounds = [s * NP for s in range(1, R)]

fig, axes = plt.subplots(1, 3, figsize=(16, 4.2))

ax = axes[0]
ax.plot(res["kl_jlens"], label="Jacobian lens", lw=1.5)
ax.plot(res["kl_logitlens"], label="logit-lens", lw=1.5)
for b in bounds:
    ax.axvline(b, color="grey", ls=":", lw=0.8)
ax.set_xlabel("virtual layer (48 layers x 4 loop steps)")
ax.set_ylabel("KL(model || lens)  [nats]")
ax.set_title("Lens quality across the unrolled loop (WikiText)")
ax.legend()

ax = axes[1]
ax.plot(ll["top1_agree_jlens"], label="Jacobian lens", lw=1.5)
ax.plot(ll["top1_agree_logitlens"], label="logit-lens", lw=1.5)
for b in bounds:
    ax.axvline(b, color="grey", ls=":", lw=0.8)
ax.set_xlabel("virtual layer")
ax.set_ylabel("top-1 agreement with final output")
ax.set_title("Readout agreement by virtual depth")
ax.legend()

ax = axes[2]
norms = [res["J_norm"][str(v)] for v in range(V - 1)]
dist = [res["J_dist_eye"][str(v)] for v in range(V - 1)]
ax.plot(norms, label="||J||/sqrt(d)", lw=1.5)
ax.plot(dist, label="||J - I||/sqrt(d)", lw=1.5)
for b in bounds:
    ax.axvline(b, color="grey", ls=":", lw=0.8)
ax.set_xlabel("virtual layer")
ax.set_title("Average Jacobian geometry")
ax.legend()

fig.tight_layout()
fig.savefig("ouro_lens_curves.png", dpi=150)

sim = torch.load("ouro_jsim.pt")["sim"].numpy()
fig2, ax = plt.subplots(figsize=(6.5, 5.5))
im = ax.imshow(sim, vmin=0, vmax=1, cmap="viridis", origin="lower")
for b in bounds:
    ax.axvline(b - 0.5, color="w", ls=":", lw=0.8)
    ax.axhline(b - 0.5, color="w", ls=":", lw=0.8)
ax.set_xlabel("virtual layer")
ax.set_ylabel("virtual layer")
ax.set_title("cos similarity of J across virtual layers")
fig2.colorbar(im)
fig2.tight_layout()
fig2.savefig("ouro_jsim_heatmap.png", dpi=150)
print("DONE_PLOTS")