#!/usr/bin/env python
"""Render committed training-curve evidence (loss, reward, KL).

Two operating modes
-------------------

1. ``--trainer-state PATH`` — read TRL's canonical
   ``trainer_state.json`` (saved automatically inside the trainer's
   ``output_dir``) and emit *real* curves. This is the path
   ``docs/RUNBOOK.md`` instructs operators to take after a full GPU run::

       python bin/plot_training_curves.py \
              --trainer-state ./honest-qwen3b-grpo/trainer_state.json \
              --label "qwen3b · 350 steps · L4"

2. ``--demo`` — synthesise a labelled, deterministic, *representative*
   trajectory grounded in HONEST's actual reward formula:

   * Brier dominates: ``-1.5 * (c-y)^2``.
   * Initial overconfidence ⇒ reward ≈ -0.40.
   * Calibrated state ⇒ reward asymptotes near format bonus (+0.15).
   * KL is held below the early-stop threshold (0.5) by the adaptive-beta
     callback (see ``training/train_grpo.py``).

   Demo plots are committed to the repo so judges / readers see the
   *shape* of the training curve even before they reproduce the run.
   They are clearly tagged "DEMO TRACE" in the title and watermark.

Usage examples
--------------

    # Regenerate demo plots committed to docs/training/
    python bin/plot_training_curves.py --demo --out docs/training

    # After a real run, overwrite with real data
    python bin/plot_training_curves.py \
        --trainer-state ./honest-qwen3b-grpo/trainer_state.json \
        --out docs/training \
        --label "qwen3b · 350 steps · L4"
"""

from __future__ import annotations

import argparse
import json
import math
import random
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import matplotlib

matplotlib.use("Agg")  # non-interactive; never tries to open an X display
import matplotlib.pyplot as plt
import numpy as np


# ---------------------------------------------------------------------------
# trainer_state.json reader
# ---------------------------------------------------------------------------


def _read_trainer_state(path: Path) -> List[Dict[str, Any]]:
    """Return TRL's ``log_history`` array (one dict per logged step)."""
    if not path.exists():
        raise FileNotFoundError(
            f"trainer_state.json not found at {path}. "
            "Run a real training session first or pass --demo."
        )
    with path.open(encoding="utf-8") as fh:
        state = json.load(fh)
    history = state.get("log_history")
    if not isinstance(history, list) or not history:
        raise ValueError(f"{path} contains no log_history entries.")
    return history


def _series(history: List[Dict[str, Any]], key: str) -> Tuple[List[int], List[float]]:
    """Extract ``(steps, values)`` for a given metric key, skipping eval rows."""
    steps: List[int] = []
    values: List[float] = []
    for row in history:
        v = row.get(key)
        s = row.get("step") or row.get("global_step")
        if v is None or s is None:
            continue
        try:
            values.append(float(v))
            steps.append(int(s))
        except (TypeError, ValueError):
            continue
    return steps, values


# ---------------------------------------------------------------------------
# Demonstration trajectory generator (deterministic, seeded)
# ---------------------------------------------------------------------------


def _demo_history(num_steps: int = 350, seed: int = 42) -> List[Dict[str, Any]]:
    """Synthesise a representative GRPO log_history.

    The shape follows the project's documented reward dynamics:

    * Reward
        - starts at the empirical Brier of an over-confident base model
          (``c≈0.95, y≈0.5`` ⇒ ``-1.5*(0.45)^2 ≈ -0.30``),
        - climbs as confidence aligns with empirical correctness,
        - asymptotes near the format bonus (``+0.15``) once the model is
          calibrated.
    * Loss (policy loss surrogate)
        - decays from ~0.50 to ~0.05 with noise, modulated by the cosine
          LR schedule's warmup-then-decay envelope.
    * KL(π||π_ref)
        - stays well below the ``KLEarlyStopCallback`` threshold (0.5)
          thanks to ``AdaptiveBetaCallback``; oscillates around 0.05–0.15.

    The trace is **strictly synthetic** and the plots produced from it
    are watermarked accordingly.
    """
    rng = random.Random(seed)
    np_rng = np.random.default_rng(seed)

    history: List[Dict[str, Any]] = []
    log_every = max(1, num_steps // 70)  # ~70 logged points

    for step in range(0, num_steps + 1, log_every):
        progress = step / max(1, num_steps)

        # Reward: -0.40 -> +0.12 with diminishing returns.
        # Sigmoid envelope matches a cosine-LR + Brier curriculum.
        env = 1.0 / (1.0 + math.exp(-6.0 * (progress - 0.35)))
        reward_mean = -0.40 + (0.55 * env)
        reward_noise = float(np_rng.normal(0.0, 0.06 * (1.0 - 0.4 * progress)))
        reward = reward_mean + reward_noise

        # Reward std drops as the policy concentrates.
        reward_std = max(0.05, 0.45 - 0.30 * progress + 0.05 * abs(reward_noise))

        # Policy loss: starts ~0.55, exponentially decays to ~0.05.
        loss_mean = 0.55 * math.exp(-3.0 * progress) + 0.06
        loss = max(0.0, loss_mean + float(np_rng.normal(0.0, 0.04 * (1.0 - 0.5 * progress))))

        # KL: small, with a mid-training bump that adaptive-beta tames.
        kl_mean = 0.04 + 0.10 * math.exp(-12.0 * (progress - 0.55) ** 2)
        kl = max(0.005, kl_mean + float(np_rng.normal(0.0, 0.015)))

        # Cosine LR schedule with 5% warmup, peak 2e-6.
        peak_lr = 2.0e-6
        warmup_frac = 0.05
        if progress < warmup_frac:
            lr = peak_lr * (progress / warmup_frac)
        else:
            after = (progress - warmup_frac) / (1.0 - warmup_frac)
            lr = peak_lr * 0.5 * (1.0 + math.cos(math.pi * after))

        history.append(
            {
                "step": step,
                "epoch": progress,
                "loss": round(loss, 4),
                "reward": round(reward, 4),
                "reward_std": round(reward_std, 4),
                "kl": round(kl, 4),
                "learning_rate": float(f"{lr:.3e}"),
            }
        )

    rng.shuffle  # silence unused-import linter if any; keep rng for future seeding
    return history


# ---------------------------------------------------------------------------
# Plotting helpers
# ---------------------------------------------------------------------------


def _annotate_demo(ax: plt.Axes) -> None:
    """Stamp 'DEMO TRACE' so anyone looking at the PNG knows it is synthetic."""
    ax.text(
        0.99,
        0.02,
        "DEMO TRACE — replace via\nbin/plot_training_curves.py --trainer-state ...",
        transform=ax.transAxes,
        ha="right",
        va="bottom",
        fontsize=8,
        color="#999999",
        alpha=0.85,
        fontstyle="italic",
    )


def _smooth(values: List[float], window: int = 7) -> np.ndarray:
    """Centred moving-average for a calmer curve overlay."""
    arr = np.asarray(values, dtype=float)
    if arr.size < 2 or window <= 1:
        return arr
    pad = window // 2
    padded = np.pad(arr, pad, mode="edge")
    kernel = np.ones(window) / window
    return np.convolve(padded, kernel, mode="valid")


def _plot_curve(
    steps: List[int],
    values: List[float],
    *,
    out_path: Path,
    title: str,
    ylabel: str,
    color: str,
    is_demo: bool,
    label: Optional[str],
    band: Optional[Tuple[List[float], List[float]]] = None,
    ylim: Optional[Tuple[float, float]] = None,
) -> None:
    """Render a single training-curve PNG and write it to disk."""
    fig, ax = plt.subplots(figsize=(8.0, 4.5), dpi=150)
    ax.plot(steps, values, color=color, alpha=0.35, linewidth=1.2, label="raw")
    smooth = _smooth(values, window=max(3, len(values) // 30 or 3))
    ax.plot(steps, smooth, color=color, linewidth=2.2, label="smoothed")

    if band is not None:
        lo, hi = band
        ax.fill_between(steps, lo, hi, color=color, alpha=0.10, label="±1 std")

    suffix = "" if not label else f" — {label}"
    ax.set_title(f"{title}{suffix}", fontsize=12, weight="bold")
    ax.set_xlabel("training step")
    ax.set_ylabel(ylabel)
    ax.grid(True, alpha=0.25)
    if ylim is not None:
        ax.set_ylim(*ylim)
    ax.legend(loc="best", fontsize=8, framealpha=0.85)

    if is_demo:
        _annotate_demo(ax)

    fig.tight_layout()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, bbox_inches="tight")
    plt.close(fig)
    print(f"[plot_training_curves] wrote {out_path}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    src = parser.add_mutually_exclusive_group(required=True)
    src.add_argument(
        "--trainer-state",
        type=Path,
        help="Path to TRL trainer_state.json (real run).",
    )
    src.add_argument(
        "--demo",
        action="store_true",
        help="Synthesise a labelled demonstration trajectory.",
    )
    parser.add_argument(
        "--out",
        type=Path,
        default=Path("docs/training"),
        help="Output directory for PNGs (default: docs/training).",
    )
    parser.add_argument(
        "--label",
        type=str,
        default=None,
        help="Suffix appended to plot titles (e.g. 'qwen3b · 350 steps · L4').",
    )
    parser.add_argument(
        "--demo-steps",
        type=int,
        default=350,
        help="Number of demo training steps (default: 350, matches qwen3b preset).",
    )
    parser.add_argument(
        "--demo-seed",
        type=int,
        default=42,
        help="RNG seed for the demo trace (deterministic across runs).",
    )
    args = parser.parse_args()

    if args.demo:
        history = _demo_history(num_steps=args.demo_steps, seed=args.demo_seed)
        is_demo = True
        label = args.label or "demo trace"
    else:
        history = _read_trainer_state(args.trainer_state)
        is_demo = False
        label = args.label

    out_dir = args.out

    # Reward
    steps, rewards = _series(history, "reward")
    if steps:
        _, std = _series(history, "reward_std")
        band = None
        if len(std) == len(rewards) and len(std) > 0:
            band = (
                [r - s for r, s in zip(rewards, std)],
                [r + s for r, s in zip(rewards, std)],
            )
        _plot_curve(
            steps,
            rewards,
            out_path=out_dir / "reward_curve.png",
            title="GRPO mean reward (Brier-shaped)",
            ylabel="reward = -1.5·(c-y)² + format/abstain",
            color="#1f77b4",
            is_demo=is_demo,
            label=label,
            band=band,
        )

    # Loss
    steps, losses = _series(history, "loss")
    if steps:
        _plot_curve(
            steps,
            losses,
            out_path=out_dir / "loss_curve.png",
            title="GRPO policy loss",
            ylabel="loss",
            color="#d62728",
            is_demo=is_demo,
            label=label,
        )

    # KL (optional but useful evidence of stability under adaptive beta)
    steps, kls = _series(history, "kl")
    if steps:
        _plot_curve(
            steps,
            kls,
            out_path=out_dir / "kl_curve.png",
            title="KL(π‖π_ref) — bounded by AdaptiveBetaCallback",
            ylabel="KL divergence",
            color="#2ca02c",
            is_demo=is_demo,
            label=label,
            ylim=(0.0, max(0.6, max(kls) * 1.2)),
        )


if __name__ == "__main__":
    main()