"""Figure: one row per beta2: left = rms_grad and rms_sqrt_vhat vs step (log y, eps line), right = test acc.
Last row: measured mean period vs 1/(1-beta2) with the law T = 2*L/(1-beta2).  Run from var_clamp/ with /work/venv-plot/bin/python."""
import csv, os, sys, numpy as np, matplotlib
matplotlib.use("Agg"); import matplotlib.pyplot as plt
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
BLUE, ORANGE, AQUA, YELLOW = "#2a78d6", "#eb6834", "#1baf7a", "#eda100"
RUNS = [(0.8, "probe_N60_b20.8_log10_eps1e-8"), (0.9, "probe_N60_b20.9_log10_eps1e-8"), (0.95, "probe_N60_b20.95_eps1e-8"),
        (0.98, "probe_N60_eps1e-8"), (0.99, "probe_N60_b20.99_eps1e-8"), (0.995, "probe_N60_b20.995_eps1e-8")]
EXTRA = [(0.95, "probe_N60_b20.95_seed1_eps1e-8", "seed 1"), (0.98, "probe_N60_eps1e-16", "eps 1e-16")]
RUNS = [(b, d) for b, d in RUNS if os.path.exists(f"{d}/probe.csv")]
def load(p):
    r = list(csv.DictReader(open(p))); return {k: np.array([float(x[k]) for x in r]) for k in r[0]}
n = len(RUNS)
fig = plt.figure(figsize=(12, 2.1*n + 3.2))
gs = fig.add_gridspec(n + 1, 2, width_ratios=[3, 1], height_ratios=[1]*n + [1.5])
for row, (b2, d) in enumerate(RUNS):
    p, lg = load(f"{d}/probe.csv"), load(f"{d}/log.csv")
    xmax = min(p['step'][-1], 24000)
    a = fig.add_subplot(gs[row, 0])
    a.semilogy(p['step'], p['rms_grad'], color=BLUE, lw=0.7, label="rms grad")
    a.semilogy(p['step'], p['rms_sqrt_vhat'], color=ORANGE, lw=1.0, label=r"rms $\sqrt{\hat v}$")
    a.axhline(1e-8, color="0.4", ls="--", lw=0.8, label="eps = 1e-8")
    a.set_xlim(0, xmax); a.set_ylim(1e-10, 1); a.grid(alpha=0.25, lw=0.5)
    a.set_title(rf"$\beta_2$ = {b2}   (1/(1-$\beta_2$) = {1/(1-b2):.0f}; logged every {int(p['step'][2]-p['step'][1])} steps)", fontsize=9, loc="left")
    a.set_ylabel("per-coord RMS", fontsize=8)
    if row == 0: a.legend(fontsize=7, loc="upper right", ncol=3, frameon=False)
    a2 = fig.add_subplot(gs[row, 1])
    a2.plot(lg['step'], lg['test_acc'], color=AQUA, lw=1.0)
    a2.set_xlim(0, xmax); a2.set_ylim(0, 0.3); a2.grid(alpha=0.25, lw=0.5); a2.set_ylabel("test acc", fontsize=8)
    for ax in (a, a2): ax.tick_params(labelsize=7)
    if row == n - 1: a.set_xlabel("step", fontsize=8); a2.set_xlabel("step", fontsize=8)
# summary panel
from period_analysis import analyse
ax = fig.add_subplot(gs[n, :])
xs, ys, es = [], [], []
for b2, d in RUNS:
    o = analyse(d, b2, verbose=False)
    if b2 < 0.85:
        ax.errorbar([1/(1-b2)], [o['per_mean']], yerr=[o['per_sd']], fmt="o", mfc="white", color=BLUE, ms=5, capsize=3, label=f"beta2 {b2} (law broken; not fitted)")
    elif np.isfinite(o['per_mean']): xs.append(1/(1-b2)); ys.append(o['per_mean']); es.append(o['per_sd'])
ax.errorbar(xs, ys, yerr=es, fmt="o", color=BLUE, ms=5, capsize=3, label="measured mean period (sd), seed 0, eps 1e-8")
for b2, d, lab in EXTRA:
    if os.path.exists(f"{d}/probe.csv"):
        o = analyse(d, b2, verbose=False)
        ax.errorbar([1/(1-b2)*(1.04 if 'seed' in lab else 0.96)], [o['per_mean']], yerr=[o['per_sd']], fmt="s" if 'seed' in lab else "D",
                    color=AQUA if 'seed' in lab else YELLOW, ms=5, capsize=3, label=f"{lab} (beta2 {b2})")
xx = np.array([3, 300])
L = np.mean([y*(1/x)/2 for x, y in zip(xs, ys)])
ax.plot(xx, 2*L*xx, color=ORANGE, lw=1, label=f"T = 2L/(1-beta2), L = {L:.1f} (mean over beta2 0.9-0.995)")
ax.set_xscale("log"); ax.set_yscale("log"); ax.grid(alpha=0.25, lw=0.5, which="both")
ax.set_xlabel(r"1/(1-$\beta_2$)", fontsize=8); ax.set_ylabel("spike period (steps)", fontsize=8); ax.tick_params(labelsize=7)
ax.legend(fontsize=7, frameon=False, loc="upper left")
fig.suptitle("N=60 norm clamp, wd 0, eps 1e-8, AdamW lr 1e-3, beta1 0.9: slingshot cycle vs beta2 (one seed per beta2 unless marked)", fontsize=10)
fig.tight_layout(); fig.savefig("fig_period_beta2.png", dpi=150); print("saved fig_period_beta2.png", [b for b, _ in RUNS], "L=", L)