import json

import matplotlib

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

res = json.load(open("ouro_targets_analysis.json"))

fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))

ax = axes[0]
for t, c in res["norm_vs_dist"].items():
    ax.plot(range(1, len(c) + 1), c, label=f"target v{t}", lw=1.5)
ax.set_xlabel("distance to target (virtual layers)")
ax.set_ylabel("||J|| / sqrt(d)")
ax.set_title("Mean-Jacobian norm vs distance, per loop-end target")
ax.legend()

ax = axes[1]
sims = res["analogous_sim_per_offset"]
ax.plot(range(len(sims)), sims, lw=1.5, label="analogous position, across loops")
ax.axhline(res["mean_mismatched_sim"], color="grey", ls="--", lw=1,
           label=f"mismatched-offset baseline ({res['mean_mismatched_sim']:.2f})")
ax.set_xlabel("physical layer offset l within loop")
ax.set_ylabel("mean cos sim of J_{s*48+l -> loop end}")
ax.set_ylim(0, 1)
ax.set_title("Weight-tying: same transport repeats each loop")
ax.legend()

fig.tight_layout()
fig.savefig("ouro_targets_curves.png", dpi=150)
print("DONE_TARGET_PLOTS")