"""Generate reward + loss curve PNGs from a TRL training run.

Reads trainer_state.json (TRL writes one per checkpoint) and creates:
  outputs/plots/reward_curve.png
  outputs/plots/loss_curve.png
  outputs/plots/grad_norm.png
  outputs/plots/training_summary.png  (combined 2x2 panel)

These are committed-to-repo PNGs that judges expect in the README.
"""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib

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


OUTPUTS = Path("outputs")
PLOTS = OUTPUTS / "plots"
PLOTS.mkdir(parents=True, exist_ok=True)


def find_trainer_state() -> Path | None:
    candidates = list(OUTPUTS.rglob("trainer_state.json"))
    if not candidates:
        return None
    candidates.sort(key=lambda p: p.stat().st_mtime, reverse=True)
    return candidates[0]


def extract_series(log_history: list[dict]) -> dict[str, list]:
    series: dict[str, list] = {"step": []}
    for entry in log_history:
        step = entry.get("step")
        if step is None:
            continue
        for k, v in entry.items():
            if k == "step":
                continue
            if not isinstance(v, (int, float)):
                continue
            series.setdefault(k, []).append((step, float(v)))
    return series


def plot_metric(series: dict, key: str, title: str, ylabel: str, outfile: Path) -> bool:
    if key not in series or len(series[key]) < 2:
        print(f"[skip] no series for {key}")
        return False
    xs, ys = zip(*series[key])
    fig, ax = plt.subplots(figsize=(8, 5), dpi=120)
    ax.plot(xs, ys, marker=".", linewidth=2, color="#1f77b4")
    ax.set_xlabel("Training step")
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(outfile)
    plt.close(fig)
    print(f"[ok]   {outfile}")
    return True


def plot_summary(series: dict, outfile: Path) -> None:
    panels = [
        ("loss", "Loss", "loss"),
        ("reward", "Mean reward", "reward"),
        ("grad_norm", "Grad norm", "grad_norm"),
        ("kl", "KL divergence", "kl"),
    ]
    fig, axes = plt.subplots(2, 2, figsize=(13, 9), dpi=120)
    for ax, (key, title, ylabel) in zip(axes.flat, panels):
        if key not in series or len(series[key]) < 2:
            ax.set_title(f"{title} (no data)")
            ax.axis("off")
            continue
        xs, ys = zip(*series[key])
        ax.plot(xs, ys, marker=".", linewidth=2)
        ax.set_xlabel("Step")
        ax.set_ylabel(ylabel)
        ax.set_title(title)
        ax.grid(True, alpha=0.3)
    fig.suptitle("FATHOM — GRPO training (Qwen 1.5B + LoRA, OpenEnv multi-turn)", fontsize=13)
    fig.tight_layout()
    fig.savefig(outfile)
    plt.close(fig)
    print(f"[ok]   {outfile}")


def main() -> int:
    state_path = find_trainer_state()
    if state_path is None:
        print("ERROR: no trainer_state.json found under outputs/")
        return 1
    print(f"Reading {state_path}")
    state = json.loads(state_path.read_text())
    series = extract_series(state.get("log_history", []))
    print(f"Series found: {sorted(series.keys())}")

    plot_metric(series, "loss", "Training loss", "loss", PLOTS / "loss_curve.png")
    plot_metric(series, "reward", "Mean reward (composite)", "reward", PLOTS / "reward_curve.png")
    plot_metric(series, "grad_norm", "Gradient norm", "grad_norm", PLOTS / "grad_norm.png")
    plot_metric(series, "kl", "KL divergence", "KL", PLOTS / "kl_curve.png")
    plot_summary(series, PLOTS / "training_summary.png")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())