#!/usr/bin/env python3
"""Recompute run summaries and the LR comparison from compact curated evidence.

Run with no arguments from any directory. The standard library produces the CSV
and JSON summaries; matplotlib is required only for the PNG chart. No trainer
state, checkpoints, model weights, or tokenizers are needed.

The gradient P95 matches the archived summary: sort N observed norms and choose
the zero-based index floor(0.95 * (N - 1)), without interpolation.
"""
import csv
import json
import math
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
RESULTS = ROOT / "results"


def read_rows(filename):
    with (RESULTS / filename).open(newline="", encoding="utf-8") as handle:
        return list(csv.DictReader(handle))


def summarize():
    evaluations = defaultdict(list)
    gradients = defaultdict(list)
    for row in read_rows("eval_loss_history.csv"):
        evaluations[row["run"]].append(row)
    for row in read_rows("gradient_history.csv"):
        gradients[row["run"]].append(float(row["grad_norm"]))
    raw_summary = {float(row["learning_rate"]): row for row in read_rows("aya_300k_lr_summary.csv")}
    summaries = []
    for run, rows in sorted(evaluations.items()):
        rows.sort(key=lambda row: int(row["step"]))
        lr = float(rows[0]["learning_rate"])
        if any(float(row["learning_rate"]) != lr for row in rows):
            raise ValueError(f"Inconsistent learning rate in {run}")
        trial = ROOT / "trials" / run
        train = json.loads((trial / "train_results.json").read_text(encoding="utf-8"))
        final_eval = json.loads((trial / "eval_results.json").read_text(encoding="utf-8"))
        norms = sorted(gradients[run])
        if not norms or not (trial / "TRAINING_COMPLETE").is_file():
            raise ValueError(f"Incomplete evidence for {run}")
        final = rows[-1]
        best = min(rows, key=lambda row: float(row["eval_loss"]))
        summary = {
            "learning_rate": lr, "run": run,
            "global_step": int(raw_summary[lr]["global_step"]),
            "epoch": float(train["epoch"]),
            "final_eval_step": int(final["step"]),
            "final_eval_loss": float(final["eval_loss"]),
            "min_eval_step": int(best["step"]),
            "min_eval_loss": float(best["eval_loss"]),
            "max_grad_norm": max(norms),
            "p95_grad_norm": norms[math.floor(0.95 * (len(norms) - 1))],
            "train_runtime": float(train["train_runtime"]),
            "eval_observations": len(rows),
            "gradient_observations": len(norms),
        }
        if summary["final_eval_loss"] != float(final_eval["eval_loss"]):
            raise ValueError(f"Final evaluation disagrees with eval_results.json: {run}")
        for field in ("global_step", "final_eval_step", "final_eval_loss", "min_eval_step",
                      "min_eval_loss", "max_grad_norm", "p95_grad_norm"):
            if float(summary[field]) != float(raw_summary[lr][field]):
                raise ValueError(f"Archived summary disagrees at {run}: {field}")
        summaries.append(summary)
    if len(summaries) != 6 or any(row["global_step"] != 18688 or row["epoch"] != 1.0
                                  or row["eval_observations"] != 19 or row["gradient_observations"] != 1868
                                  for row in summaries):
        raise ValueError("Expected six completed one-epoch runs with 19 eval and 1,868 gradient observations each")
    with (RESULTS / "training_run_summary.csv").open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(summaries[0]))
        writer.writeheader()
        writer.writerows(summaries)
    (RESULTS / "training_run_summary.json").write_text(json.dumps(summaries, indent=2) + "\n", encoding="utf-8")
    return summaries


def plot(summaries):
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise SystemExit("Summaries recomputed. Install matplotlib to regenerate the PNG chart.") from exc
    plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 11})
    fig, ax = plt.subplots(figsize=(8, 4.8), layout="constrained")
    x = [row["learning_rate"] * 1e4 for row in summaries]
    y = [row["final_eval_loss"] for row in summaries]
    ax.plot(x, y, color="#335c81", marker="o", linewidth=1.6, markersize=6)
    selected = min(summaries, key=lambda row: row["final_eval_loss"])
    sx, sy = selected["learning_rate"] * 1e4, selected["final_eval_loss"]
    ax.scatter([sx], [sy], color="#ba4a00", s=75, zorder=3)
    for xi, yi in zip(x, y):
        ax.annotate(f"{yi:.5f}", (xi, yi), xytext=(0, 10), textcoords="offset points", ha="center", fontsize=9)
    ax.annotate("Selected LR = 0.00010", (sx, sy), xytext=(20, -24),
                textcoords="offset points", color="#8a3700", fontsize=10)
    ax.set(title="Aya 300k: final validation loss by learning rate",
           xlabel="Learning rate (×10⁻⁴)", ylabel="Final validation loss (lower is better)")
    ax.set_xticks(x, [f"{value:.1f}" for value in x])
    ax.set_ylim(min(y) - 0.009, max(y) + 0.012)
    ax.grid(axis="y", color="#d8dee5", linewidth=0.7)
    ax.spines[["top", "right"]].set_visible(False)
    fig.savefig(RESULTS / "lr_vs_final_eval_loss.png", dpi=200, metadata={"Description": "Six manually chosen learning rates; final evaluation at step 18688 after one epoch."})
    plt.close(fig)


def main():
    summaries = summarize()
    plot(summaries)
    best = min(summaries, key=lambda row: row["final_eval_loss"])
    print(f"Verified {len(summaries)} runs. Selected {best['run']}: final eval loss {best['final_eval_loss']:.15g}")
    print(f"Saved summaries and chart in {RESULTS}")


if __name__ == "__main__":
    main()