#!/usr/bin/env python3
"""Appendix training-dynamics figures (4B / 8B), one single-figure PDF+PNG per
panel — composed into a 2x2 layout in LaTeX via \\subfigure. Follows doc/vis.md:
restricted academic palette, sans-serif, large fonts, tight crop, B&W-safe
(solid/dashed + circle/square markers), "Ours" (GroundFlow) gets the prominent
coral colour consistently.

Outputs (to doc/figures/ and figures/):
  figA_loss_4b.{pdf,png}  figA_loss_8b.{pdf,png}
  figA_gn_4b.{pdf,png}    figA_gn_8b.{pdf,png}
"""
import json, math, os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker

M = "/mnt/bn/leonworkspace/terry/model"
TRAINER = {
    "4B": {
        "SD-RPN":     f"{M}/qwen3vl-4b-roi-K24T3-185k-ddp-verify/trainer_state.json",
        "GroundFlow": f"{M}/qwen3vl-4b-roi-K24T3-185k-selfdistill-run7b-sdpa-iwa/trainer_state.json",
    },
    "8B": {
        "SD-RPN":     f"{M}/qwen3vl-8b-roi-K24T3-185k-sdrpn/trainer_state.json",
        "GroundFlow": f"{M}/qwen3vl-8b-roi-K24T3-185k-a0.667-t1.5-iwa-g1.0/trainer_state.json",
    },
}
# doc/vis.md palette: GroundFlow (=Ours) = coral #f57c6e (prominent, consistent).
# Per request, only the GroundFlow curve is plotted (SD-RPN's loss is not shown).
STYLE = {
    "GroundFlow": dict(color="#f57c6e", ls="-",  marker="o", label="GroundFlow"),
}
ORDER = ["GroundFlow"]
OUTDIRS = ["/opt/tiger/thothvl_pretrain/doc/figures", "/opt/tiger/thothvl_pretrain/figures"]

def load(path):
    lh = json.load(open(path))["log_history"]
    rows = [(d["step"], d["loss"], d.get("grad_norm")) for d in lh if "loss" in d]
    step = np.array([r[0] for r in rows], float)
    loss = np.array([r[1] for r in rows], float)
    gn   = np.array([(r[2] if r[2] is not None else np.nan) for r in rows], float)
    fin  = [d for d in lh if "train_loss" in d]
    return step, loss, gn, (fin[-1]["train_loss"] if fin else float(loss[-1]))

def ema(x, alpha=0.06):
    y = np.empty_like(x); m = next(v for v in x if math.isfinite(v))
    for i, v in enumerate(x):
        if math.isfinite(v):
            m = alpha * v + (1 - alpha) * 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": 13, "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,
})

def make_panel(model, kind, fname):
    """kind in {'loss','gn'}"""
    fig, ax = plt.subplots(figsize=(4.0, 3.0))
    for name in ORDER:
        st, ls, gn, fl = load(TRAINER[model][name])
        y = ls if kind == "loss" else gn
        s = STYLE[name]
        idx = np.linspace(0, len(st) - 1, 8).round().astype(int)  # 8 markers, B&W aid
        ax.plot(st, y, color=s["color"], lw=0.6, alpha=0.16)               # raw, faint
        ax.plot(st, ema(y), color=s["color"], ls=s["ls"], lw=2.0,          # EMA
                marker=s["marker"], markevery=list(idx), ms=4.5, mew=0,
                label=s["label"], zorder=3 if name == "GroundFlow" else 2)
        if kind == "loss":
            yi = ema(y)[-1]
            ax.annotate(f"{fl:.2f}" if fl >= 0.1 else f"{fl:.3f}",
                        xy=(st[-1], yi), xytext=(5, 0), textcoords="offset points",
                        va="center", fontsize=10, color=s["color"], fontweight="bold")
    if kind == "loss":
        ax.set_ylabel("training loss")
        ax.set_ylim(bottom=0)                         # linear: clean single-curve shape
    else:
        ax.set_yscale("log"); ax.set_ylabel("gradient norm")  # spans ~2 decades
    ax.set_xlabel("optimizer step")
    ax.set_title(f"Qwen3-VL-{model}")
    ax.set_xlim(-30, 1480)
    ax.xaxis.set_major_locator(mticker.MultipleLocator(500))
    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 OUTDIRS:
        os.makedirs(d, exist_ok=True)
        for ext in ("pdf", "png"):
            fig.savefig(f"{d}/{fname}.{ext}", dpi=300)
    plt.close(fig)
    print(f"  wrote {fname}.{{pdf,png}}")

for model in ("4B", "8B"):
    make_panel(model, "loss", f"figA_loss_{model.lower()}")
    make_panel(model, "gn",   f"figA_gn_{model.lower()}")
print("done.")