#!/usr/bin/env python3
"""Benchmark graphs for the KOLM-Alpha release: KOLM-Alpha vs TMT."""
import csv
import os
import matplotlib.pyplot as plt

os.makedirs("assets", exist_ok=True)
KOLM = "#4C6FFF"   # KOLM-Alpha
TMT = "#FF7A45"    # Transformer Model Twin


def load(path):
    toks, val = [], []
    with open(path) as f:
        for step, t, v in csv.reader(f):
            toks.append(int(t) / 1e6)
            val.append(float(v))
    return toks, val


kt, kv = load("curve_kolm.csv")
tt, tv = load("curve_transformer.csv")

# --- Graph 1: validation-loss curves (the crossover) ---
fig, ax = plt.subplots(figsize=(8, 5), dpi=140)
ax.plot(tt, tv, color=TMT, lw=2.2, label="TMT (Transformer twin, 17.5M)")
ax.plot(kt, kv, color=KOLM, lw=2.2, label="KOLM-Alpha (16.9M)")
ax.axhline(kv[-1], color=KOLM, ls=":", lw=1, alpha=0.6)
ix = next(i for i in range(1, len(kv)) if kv[i] <= tv[i])
ax.annotate("crossover", xy=(kt[ix], kv[ix]), xytext=(kt[ix] + 1.5, kv[ix] + 0.25),
            arrowprops=dict(arrowstyle="->", color="#555"), fontsize=10, color="#555")
ax.set_xlabel("training tokens (millions)")
ax.set_ylabel("validation loss")
ax.set_title("KOLM-Alpha vs TMT — Validation Loss (TinyStories)", fontweight="bold")
ax.legend(frameon=False)
ax.grid(alpha=0.15)
for s in ("top", "right"):
    ax.spines[s].set_visible(False)
fig.tight_layout()
fig.savefig("assets/val_loss_curve.png")

# --- Graph 2: final perplexity + loss bars ---
fig, axes = plt.subplots(1, 2, figsize=(8, 4.2), dpi=140)
for ax, vals, title in [
    (axes[0], (14.80, 15.16), "Final validation perplexity"),
    (axes[1], (2.6946, 2.7188), "Final validation loss"),
]:
    bars = ax.bar(["KOLM-Alpha", "TMT"], vals, color=[KOLM, TMT], width=0.6)
    ax.set_title(title, fontweight="bold", fontsize=11)
    ax.set_ylim(0, max(vals) * 1.15)
    for b, v in zip(bars, vals):
        ax.text(b.get_x() + b.get_width() / 2, v, f"{v:.4g}",
                ha="center", va="bottom", fontsize=10, fontweight="bold")
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    ax.grid(axis="y", alpha=0.15)
fig.suptitle("KOLM-Alpha wins at 3.5% fewer parameters (16.9M vs 17.5M)",
             fontsize=11)
fig.tight_layout()
fig.savefig("assets/final_scores.png")
print("wrote assets/val_loss_curve.png, assets/final_scores.png")