"""Plot the training and validation loss curves recorded in the lab notebook.

Purpose: turn the per-epoch numbers that train-mnist.py and train-mnist-mlx.py appended to
         labbook.md into one picture per run, so the shapes from the generalisation lesson
         can be seen rather than read off a column of digits.
Platform: all (no accelerator needed; matplotlib only)
Minimum memory: 8 GB
Assumes: matplotlib is installed in the active environment, and labbook.md holds at least
         one JSON line whose "lab" field starts with "part-01/train-mnist".

Usage: python plot-curves.py [--labbook labbook.md] [--out curves.png] [--last N]

Each run becomes one panel: training loss (solid), validation loss (dashed) and a dot on
the epoch whose checkpoint was kept. The y axis is logarithmic so a run that reached 0.01
and a run stuck at 2.3 are both readable.
"""
import argparse
import json
import sys
from pathlib import Path

try:
    import matplotlib.pyplot as plt
except ImportError:
    sys.exit("plot-curves: matplotlib is not installed; run: uv pip install matplotlib "
             "(pip install matplotlib inside the NGC container)")


def load_runs(labbook: Path) -> list[dict]:
    runs = []
    for line in labbook.read_text(encoding="utf-8").splitlines():
        if not line.startswith("{"):
            continue
        try:
            record = json.loads(line)
        except json.JSONDecodeError:
            continue
        if str(record.get("lab", "")).startswith("part-01/train-mnist") and record.get("history"):
            runs.append(record)
    return runs


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--labbook", default="labbook.md")
    parser.add_argument("--out", default="curves.png")
    parser.add_argument("--last", type=int, default=0, help="plot only the last N runs (0 = all)")
    args = parser.parse_args()

    labbook = Path(args.labbook)
    if not labbook.exists():
        sys.exit(f"plot-curves: {labbook} does not exist; run train-mnist.py with --labbook first")
    runs = load_runs(labbook)
    if not runs:
        sys.exit(f"plot-curves: no part-01/train-mnist records with a history in {labbook}")
    if args.last:
        runs = runs[-args.last:]

    cols = 2 if len(runs) > 1 else 1
    rows = (len(runs) + cols - 1) // cols
    fig, axes = plt.subplots(rows, cols, figsize=(5.5 * cols, 3.6 * rows), squeeze=False)
    for ax in axes.flat[len(runs):]:
        ax.set_visible(False)

    print(f"{'run':<4} {'label':<16} {'lr':>8} {'epochs':>6} {'train':>6} {'best':>4} {'min val loss':>12} {'final train':>11}")
    for i, (run, ax) in enumerate(zip(runs, axes.flat), start=1):
        epochs = [h["epoch"] for h in run["history"]]
        train = [h["train_loss"] for h in run["history"]]
        val = [h["val_loss"] for h in run["history"]]
        best = run.get("best_epoch") or 0
        ax.plot(epochs, train, marker=".", label="training loss")
        ax.plot(epochs, val, marker=".", linestyle="--", label="validation loss")
        if best:
            ax.plot([best], [val[best - 1]], marker="o", markersize=9, linestyle="none",
                    label=f"kept checkpoint (epoch {best})")
        ax.set_yscale("log")
        ax.set_xlabel("epoch")
        ax.set_ylabel("cross-entropy loss")
        label = run.get("label") or f"run {i}"
        ax.set_title(f"{label}: lr {run['lr']}, {run.get('train_size', 50000):,} images, {run['device']}", fontsize=9)
        ax.grid(True, which="both", alpha=0.3)
        ax.legend(fontsize=8)
        min_val = min(v for v in val if v == v) if any(v == v for v in val) else float("nan")
        print(f"{i:<4} {label[:16]:<16} {run['lr']:>8} {run['epochs']:>6} {run.get('train_size', 50000):>6} "
              f"{best:>4} {min_val:>12.4f} {train[-1]:>11.4f}")

    fig.tight_layout()
    fig.savefig(args.out, dpi=120)
    print(f"wrote {args.out} ({len(runs)} run(s) from {labbook})")


if __name__ == "__main__":
    main()