#!/usr/bin/env python3
from __future__ import annotations

import os
from pathlib import Path

os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import matplotlib.pyplot as plt


OUT_DIR = Path("figures")
OUT_DIR.mkdir(exist_ok=True)

FINAL_LONG_RUN_PPL = 19.7822

EXPERIMENTS = [
    {
        "name": "Naive\nbaseline",
        "legend": "baseline",
        "color": "#6b7280",
        "public_ppl": 42.2650,
        "train_loss_4500": 3.6290,
        "internal_val_loss_4500": 3.6823,
        "curve": [
            (0, 10.9696), (250, 6.0841), (500, 5.2423), (750, 4.6500),
            (1000, 4.3752), (1250, 4.2232), (1500, 4.1024),
            (1750, 3.9847), (2000, 3.9412), (2250, 3.8667),
            (2500, 3.8559), (2750, 3.7679), (3000, 3.7503),
            (3250, 3.7585), (3500, 3.7046), (3750, 3.7635),
            (4000, 3.6673), (4250, 3.6389), (4500, 3.6823),
        ],
    },
    {
        "name": "Mixed\ndata",
        "legend": "data",
        "color": "#2563eb",
        "public_ppl": 38.7702,
        "train_loss_4500": 3.4646,
        "internal_val_loss_4500": 3.4312,
        "curve": [
            (0, 10.9799), (250, 5.8726), (500, 5.1043), (750, 4.5490),
            (1000, 4.1688), (1250, 4.0293), (1500, 3.9846),
            (1750, 3.8147), (2000, 3.7881), (2250, 3.7224),
            (2500, 3.6965), (2750, 3.6746), (3000, 3.5739),
            (3250, 3.5122), (3500, 3.5903), (3750, 3.4440),
            (4000, 3.5392), (4250, 3.4584), (4500, 3.4312),
        ],
    },
    {
        "name": "Muon\noptimizer",
        "legend": "Muon",
        "color": "#dc2626",
        "public_ppl": 40.0987,
        "train_loss_4500": 3.5600,
        "internal_val_loss_4500": 3.6192,
        "curve": [
            (0, 10.9696), (250, 5.9106), (500, 5.1345), (750, 4.4705),
            (1000, 4.2422), (1250, 4.1028), (1500, 3.9927),
            (1750, 3.8881), (2000, 3.8523), (2250, 3.7771),
            (2500, 3.7735), (2750, 3.6887), (3000, 3.6751),
            (3250, 3.6859), (3500, 3.6309), (3750, 3.6952),
            (4000, 3.6025), (4250, 3.5744), (4500, 3.6192),
        ],
    },
    {
        "name": "Lyra\narchitecture",
        "legend": "arch",
        "color": "#059669",
        "public_ppl": 36.5445,
        "train_loss_4500": 3.4135,
        "internal_val_loss_4500": 3.4889,
        "curve": [
            (0, 10.9630), (250, 5.5041), (500, 4.6324), (750, 4.2915),
            (1000, 4.0727), (1250, 3.9969), (1500, 3.8429),
            (1750, 3.7813), (2000, 3.7209), (2250, 3.7253),
            (2500, 3.6427), (2750, 3.6593), (3000, 3.5319),
            (3250, 3.6220), (3500, 3.5259), (3750, 3.5657),
            (4000, 3.5040), (4250, 3.5018), (4500, 3.4889),
        ],
    },
    {
        "name": "Combined\nshort run",
        "legend": "combined",
        "color": "#7c3aed",
        "public_ppl": 32.1195,
        "train_loss_4500": 3.2819,
        "internal_val_loss_4500": 3.3663,
        "curve": [
            (0, 10.9511), (250, 5.5112), (500, 4.4614), (750, 4.0177),
            (1000, 3.8477), (1250, 3.7237), (1500, 3.7700),
            (1750, 3.6963), (2000, 3.5705), (2250, 3.5041),
            (2500, 3.4608), (2750, 3.3681), (3000, 3.3872),
            (3250, 3.4139), (3500, 3.3327), (3750, 3.3329),
            (4000, 3.3754), (4250, 3.2846), (4500, 3.3663),
        ],
    },
]


def main() -> None:
    plt.style.use("seaborn-v0_8-whitegrid")
    fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2), constrained_layout=True)

    ax = axes[0]
    labels = [exp["name"] for exp in EXPERIMENTS]
    ppls = [exp["public_ppl"] for exp in EXPERIMENTS]
    colors = [exp["color"] for exp in EXPERIMENTS]
    bars = ax.bar(labels, ppls, color=colors, width=0.68)
    ax.axhline(FINAL_LONG_RUN_PPL, color="#111827", linewidth=1.8, linestyle="--")
    ax.text(
        0.02,
        FINAL_LONG_RUN_PPL + 0.3,
        f"final long run: {FINAL_LONG_RUN_PPL:.2f}",
        transform=ax.get_yaxis_transform(),
        ha="left",
        va="bottom",
        fontsize=9,
        color="#111827",
    )
    ax.set_title("Course Public Validation Perplexity", fontsize=13, weight="bold")
    ax.set_ylabel("perplexity, lower is better")
    ax.set_ylim(FINAL_LONG_RUN_PPL * 0.85, max(ppls) * 1.12)
    for bar, val in zip(bars, ppls):
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            val,
            f"{val:.1f}",
            ha="center",
            va="bottom",
            fontsize=9,
        )

    ax = axes[1]
    for exp in EXPERIMENTS:
        steps = [x for x, _ in exp["curve"]]
        losses = [y for _, y in exp["curve"]]
        ax.plot(
            steps,
            losses,
            marker="o",
            linewidth=2.0,
            markersize=4,
            color=exp["color"],
            label=exp["legend"],
        )
    ax.set_title("Short-Run Heldout Loss Curves", fontsize=13, weight="bold")
    ax.set_xlabel("training iteration")
    ax.set_ylabel("validation loss")
    ax.legend(frameon=True, fontsize=9)
    ax.text(
        0.02,
        -0.18,
        "Each run changes one variable and uses the same 4,500-iteration budget; public PPL is the comparable metric.",
        transform=ax.transAxes,
        fontsize=8.5,
        color="#4b5563",
    )

    fig.suptitle("Ablation Summary for the Presentation", fontsize=15, weight="bold")
    for suffix in ("png", "pdf"):
        out = OUT_DIR / f"presentation_ablation_summary_standalone.{suffix}"
        fig.savefig(out, dpi=220)
        print(out)


if __name__ == "__main__":
    main()