#!/usr/bin/env python3
"""Chart measured training throughput from canonical aggregate JSON."""

from __future__ import annotations

import argparse
from pathlib import Path

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

from common import read_json


parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--report", type=Path, required=True)
parser.add_argument("--out", type=Path, required=True)
args = parser.parse_args()
report = read_json(args.report)
training = report.get("training") or {}
series = training.get("series") or []

fig, ax = plt.subplots(figsize=(9, 4.8))
if series:
    ax.plot([row["step"] for row in series], [row["tokens_per_second"] for row in series], color="#2563eb", lw=1.3)
    measured = training.get("mean_tokens_per_second")
    if measured is not None:
        ax.axhline(measured, color="#059669", ls="--", label=f"post-warmup mean: {measured:,.0f} tok/s")
        ax.legend()
else:
    ax.text(0.5, 0.5, "not measured", transform=ax.transAxes, ha="center", va="center", fontsize=18, color="#64748b")
ax.set(title="v25 SFT training throughput", xlabel="training step", ylabel="tokens / second")
ax.grid(alpha=0.25)
fig.tight_layout()
args.out.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(args.out, dpi=160, bbox_inches="tight")
plt.close(fig)
print(args.out)