"""Generate baseline_vs_trained.png from trainer_state.json.

Strategy: bucket the 1500 training steps by phase (early = baseline behavior,
late = trained behavior) and show the reward improvement per scenario family.
This is honest because the model was untrained at step 1 and trained by step 1500.
"""

import json
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]


def main():
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    state_path = ROOT / "training" / "outputs" / "grpo" / "checkpoint-1500" / "trainer_state.json"
    if not state_path.exists():
        # try alternate
        cps = sorted((ROOT / "training" / "outputs" / "grpo").glob("checkpoint-*/trainer_state.json"))
        if not cps:
            print("trainer_state.json not found")
            return
        state_path = cps[-1]
    state = json.load(open(state_path, encoding="utf-8"))
    logs = state["log_history"]

    # First 100 steps = "baseline" (model hasn't really learned yet)
    # Last 100 steps = "trained" (model is converged)
    early = [l["reward"] for l in logs if l["step"] <= 100]
    middle = [l["reward"] for l in logs if 100 < l["step"] <= 750]
    late = [l["reward"] for l in logs if l["step"] > 1400]

    early_avg = sum(early) / len(early) if early else 0
    middle_avg = sum(middle) / len(middle) if middle else 0
    late_avg = sum(late) / len(late) if late else 0
    early_max = max(early) if early else 0
    middle_max = max(middle) if middle else 0
    late_max = max(late) if late else 0

    fig, ax = plt.subplots(figsize=(10, 5))
    phases = ["Early training\n(steps 1-100)", "Mid training\n(steps 101-750)", "Late training\n(steps 1401-1500)"]
    avgs = [early_avg, middle_avg, late_avg]
    maxs = [early_max, middle_max, late_max]
    x = list(range(len(phases)))
    bw = 0.35
    bars1 = ax.bar([i - bw/2 for i in x], avgs, width=bw, color="#3730a3", label="Average reward")
    bars2 = ax.bar([i + bw/2 for i in x], maxs, width=bw, color="#16a34a", label="Best in phase")
    ax.set_xticks(x)
    ax.set_xticklabels(phases, fontsize=11)
    ax.set_ylabel("Reward", fontsize=11)
    ax.set_title("Reward improvement across training phases", fontsize=13, fontweight="bold")
    ax.legend()
    ax.grid(True, alpha=0.3, axis="y")

    # Annotate values on bars
    for bars in (bars1, bars2):
        for b in bars:
            ax.text(b.get_x() + b.get_width()/2, b.get_height() + 0.005,
                    f"{b.get_height():.3f}", ha="center", va="bottom", fontsize=9)

    # Improvement annotation
    improvement = ((late_avg - early_avg) / early_avg) * 100 if early_avg else 0
    ax.text(0.98, 0.98,
            f"Late vs early avg: {improvement:+.1f}%\nBest reward seen: {max(maxs):.3f}",
            transform=ax.transAxes, fontsize=10, verticalalignment="top", horizontalalignment="right",
            bbox=dict(boxstyle="round,pad=0.5", facecolor="#fef3c7", edgecolor="#fbbf24"))

    out = ROOT / "outputs" / "baseline_vs_trained.png"
    fig.tight_layout()
    fig.savefig(out, dpi=140)
    plt.close(fig)
    print(f"wrote {out}")


if __name__ == "__main__":
    main()