#!/usr/bin/env python3
"""32B α=0.667 τ=1.5: training-loss comparison — learnable IWA (γ fp32, γ_init=0,
run9, trial 300538223, in progress) vs frozen IWA (γ stuck at 1.0, the released
…iwa-0505_0321). Follows doc/vis.md style."""
import json, math, re, os
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker

# --- frozen-IWA (released) — full 1173 steps from trainer_state.json ---
TS = "/mnt/bn/leonworkspace/terry/model/qwen3vl-32b-dense-roi-K49T3-150k-confluent-a0.667-t1.5-iwa-0505_0321/trainer_state.json"
lh = json.load(open(TS))["log_history"]
fr = [(d["step"], d["loss"]) for d in lh if "loss" in d]
fr_step = np.array([x[0] for x in fr], float); fr_loss = np.array([x[1] for x in fr], float)
fr_final = [d for d in lh if "train_loss" in d][-1]["train_loss"]

# --- learnable-IWA (run9) — partial, parsed from the live stdout ---
pat = re.compile(r"\{'loss': ([0-9.]+), 'grad_norm'")
ln = [float(pat.search(l).group(1)) for l in open("/tmp/r9_loss_lines.txt") if pat.search(l)]
ln_step = np.arange(1, len(ln) + 1, dtype=float); ln_loss = np.array(ln, float)

def ema(x, a=0.06):
    y = np.empty_like(x); m = x[0]
    for i, v in enumerate(x):
        m = a * v + (1 - a) * m; y[i] = m
    return y

plt.rcParams.update({
    "font.family": "sans-serif", "font.sans-serif": ["Helvetica", "Arial", "DejaVu Sans"],
    "mathtext.fontset": "dejavusans", "font.size": 11, "axes.titlesize": 12.5,
    "axes.labelsize": 12, "xtick.labelsize": 10.5, "ytick.labelsize": 10.5, "legend.fontsize": 10,
    "axes.linewidth": 0.9, "xtick.direction": "in", "ytick.direction": "in",
    "xtick.major.size": 3.2, "ytick.major.size": 3.2, "legend.frameon": False,
    "axes.spines.top": False, "axes.spines.right": False, "savefig.bbox": "tight", "savefig.pad_inches": 0.02,
})
C_LEARN = "#f57c6e"   # Ours = learnable IWA — prominent coral
C_FROZEN = "#71b8ed"  # frozen γ=1.0 (released)

fig, ax = plt.subplots(figsize=(5.2, 3.4))
mk = lambda st: list(np.linspace(0, len(st) - 1, 8).round().astype(int))
ax.plot(fr_step, fr_loss, color=C_FROZEN, lw=0.6, alpha=0.16)
ax.plot(fr_step, ema(fr_loss), color=C_FROZEN, ls="--", lw=2.0, marker="s", markevery=mk(fr_step),
        ms=4.5, mew=0, label=r"frozen IWA ($\gamma{\equiv}1.0$, released)", zorder=2)
ax.plot(ln_step, ln_loss, color=C_LEARN, lw=0.6, alpha=0.16)
ax.plot(ln_step, ema(ln_loss), color=C_LEARN, ls="-", lw=2.2, marker="o", markevery=mk(ln_step),
        ms=4.7, mew=0, label=r"learnable IWA ($\gamma_0{=}0$, fp32) — Ours", zorder=3)
ax.annotate(f"{fr_loss[-1]:.3f}", xy=(fr_step[-1], ema(fr_loss)[-1]), xytext=(5, -4),
            textcoords="offset points", fontsize=9.5, color=C_FROZEN, fontweight="bold")
ax.annotate(f"{ln_loss[-1]:.3f}\n(@step {len(ln)}, ~{100*len(ln)//1173}%)", xy=(ln_step[-1], ema(ln_loss)[-1]),
            xytext=(6, 8), textcoords="offset points", fontsize=9.0, color=C_LEARN, fontweight="bold", va="bottom")
ax.set_yscale("log"); ax.set_xlim(-30, 1230); ax.set_ylim(0.28, 3.4)
ax.xaxis.set_major_locator(mticker.MultipleLocator(300))
ax.set_xlabel("optimizer step"); ax.set_ylabel("training loss")
ax.set_title(r"Qwen3-VL-32B, $\alpha{=}0.667$, $\tau{=}1.5$: loss floor unchanged ($\approx0.34$)")
ax.grid(True, which="both", ls=":", lw=0.5, alpha=0.35)
ax.legend(loc="upper right", handlelength=2.4, borderaxespad=0.4)
fig.tight_layout()
for d in ("/tmp", "/opt/tiger/thothvl_pretrain/figures", "/opt/tiger/thothvl_pretrain/doc/figures"):
    os.makedirs(d, exist_ok=True)
    for ext in ("pdf", "png"):
        fig.savefig(f"{d}/figA_iwa_learn_vs_frozen_a0667.{ext}", dpi=300)
print("wrote figA_iwa_learn_vs_frozen_a0667.{pdf,png}")
# numeric comparison at matched steps
print("\n  step   frozen-IWA   learnable-IWA (run9)")
for s in (1, 50, 100, 200, 300, 500, 768, 1000, 1173):
    fv = fr_loss[s-1] if s <= len(fr_loss) else None
    lv = ln_loss[s-1] if s <= len(ln_loss) else None
    print(f"  {s:>4}   {('%.4f'%fv) if fv is not None else '   -  ':>8}     {('%.4f'%lv) if lv is not None else '   - (not yet)'}")
print(f"\n  frozen final train_loss = {fr_final:.4f} (last step {fr_loss[-1]:.4f}) ; learnable still running (~{100*len(ln)//1173}%, last {ln_loss[-1]:.4f}).")