#!/usr/bin/env python3
"""Chart for reports/d58_scaling.md and the model card: KL against training tokens, one panel per codebook format
(small multiples: the formats' KLs differ 6x, so each panel has its own scale, starting at zero), from
reports/d58_curve.csv. Dashed line: llama.cpp's own file of that format (the start). Filled markers: finished files
(with an LR decay); hollow: constant-LR checkpoints; line: the finished-point fit (IQ3_S has none: it saturates by
~20M). Needs matplotlib (gen's system python3 has it).

    python3 scripts/d58_scaling_chart.py --out reports/d58_scaling.png
"""
from __future__ import annotations

import argparse
import csv
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

REPO = Path(__file__).resolve().parents[1]
SURFACE, INK, INK2, GRID = "#fcfcfb", "#0b0b0b", "#52514e", "#e6e5e1"
SERIES = [("IQ2_XXS", "2-bit IQ2_XXS", "#2a78d6", (0.0775, 0.1354, 0.637)),  # validated slots 1-3 (dataviz palette)
          ("IQ2_S", "2-bit IQ2_S", "#eb6834", (0.0506, 0.0591, 0.622)),
          ("IQ3_S", "3-bit IQ3_S", "#1baf7a", None)]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", type=Path, default=REPO / "reports" / "d58_curve.csv")
    ap.add_argument("--out", type=Path, required=True)
    a = ap.parse_args()
    rows = list(csv.DictReader(a.csv.open()))
    fig, axes = plt.subplots(1, 3, figsize=(12, 4.6), dpi=150)
    fig.patch.set_facecolor(SURFACE)
    for ax, (fmt, title, color, fit) in zip(axes, SERIES):
        ax.set_facecolor(SURFACE)
        start = next(float(r["kl"]) for r in rows if r["format"] == fmt and r["phase"] == "start")
        fin = sorted((int(r["tokens"]) / 1e6, float(r["kl"])) for r in rows
                     if r["format"] == fmt and r["phase"] in ("branch", "final"))
        con = sorted((int(r["tokens"]) / 1e6, float(r["kl"])) for r in rows
                     if r["format"] == fmt and r["phase"] == "constant")
        ax.axhline(start, color=INK2, lw=1.2, ls=(0, (4, 3)), zorder=1)
        ax.text(52, start, f"llama.cpp's own file: {start:.3f}", color=INK2, fontsize=8.5, va="bottom", ha="right")
        if fit:
            c, k, b = fit
            xs = [x / 2 for x in range(10, 101)]
            ax.plot(xs, [c + k * x ** -b for x in xs], color=color, lw=2, zorder=2)
        else:
            ax.plot([x for x, _ in fin], [y for _, y in fin], color=color, lw=2, zorder=2)
        ax.plot([x for x, _ in con], [y for _, y in con], marker="o", ms=8, mfc=SURFACE, mec=color, mew=2,
                ls="none", zorder=3)
        ax.plot([x for x, _ in fin], [y for _, y in fin], marker="o", ms=8, color=color, mec=SURFACE, mew=2,
                ls="none", zorder=4)
        best = min(fin, key=lambda t: t[1])
        ax.annotate(f"best {best[1]:.4f}\n({best[0]:.0f}M tokens, {100 * (1 - best[1] / start):.0f}% below llama.cpp)",
                    best, xytext=(6 if best[0] > 30 else 0, -36), textcoords="offset points",
                    ha="right" if best[0] > 30 else "center", fontsize=8.5, color=INK)
        ax.set_ylim(0, start * 1.12)
        ax.set_xlim(0, 53)
        ax.set_title(title, color=INK, loc="left", fontsize=11)
        ax.set_xlabel("Training tokens (millions)", color=INK2, fontsize=9)
        ax.grid(True, color=GRID, lw=0.8)
        ax.set_axisbelow(True)
        for s in ("top", "right"):
            ax.spines[s].set_visible(False)
        for s in ("left", "bottom"):
            ax.spines[s].set_color(GRID)
        ax.tick_params(colors=INK2, labelsize=8.5)
    axes[0].set_ylabel("KL to the original (lower is better)", color=INK2, fontsize=9)
    fig.suptitle("Training the codebook scales: KL against training tokens", color=INK, x=0.01, y=0.985, ha="left", fontsize=12)
    fig.text(0.01, 0.855, "Training curves, measured before attention and embeddings were converted to Q8_0 "
             "(the shipped files are about 0.001 higher). 40 × 4,096 held-out tokens.\n"
             "Filled: runs ending with an LR decay. Hollow: constant-LR checkpoints. Line: fit. Each panel has its own "
             "scale.", color=INK2, fontsize=8.5, va="bottom")
    fig.tight_layout(rect=(0, 0, 1, 0.83))
    a.out.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(a.out, facecolor=SURFACE)
    print(f"wrote {a.out}")


if __name__ == "__main__":
    main()