"""
Make plots from a run's results (JSON) — used in the README.

  python plot_results.py runs/colab-l4-2026-10-08
"""
from __future__ import annotations

import json
import os
import sys

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

SURFACE = "#fcfcfb"
INK = "#0b0b0b"
INK_2 = "#52514e"
GRID = "#e4e3df"
SERIES = {"base": "#2a78d6", "sft": "#eb6834", "dpo": "#1baf7a"}  # fixed categorical order


def _style(ax, title: str):
    ax.set_facecolor(SURFACE)
    ax.set_title(title, loc="left", color=INK, fontsize=12, fontweight="bold", pad=12)
    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=INK_2, labelsize=9)
    ax.grid(color=GRID, linewidth=0.8, which="major")
    ax.set_axisbelow(True)


def plot_pretrain(run: str):
    d = json.load(open(os.path.join(run, "pretrain.json")))
    tok = np.array([h["tokens"] for h in d["history"]]) / 1e6
    ppl = np.array([h["val_ppl"] for h in d["history"]])
    fig, ax = plt.subplots(figsize=(7, 4), facecolor=SURFACE)
    _style(ax, "Pretraining — validation perplexity")
    ax.plot(tok, ppl, color=SERIES["base"], linewidth=2)
    ax.scatter(tok, ppl, s=28, color=SERIES["base"], edgecolor=SURFACE, linewidth=2, zorder=3)
    ax.annotate(f"{ppl[-1]:.2f}", (tok[-1], ppl[-1]), xytext=(6, 6), textcoords="offset points", color=INK, fontsize=9)
    ax.annotate(f"{ppl[0]:.2f}", (tok[0], ppl[0]), xytext=(6, 0), textcoords="offset points", color=INK, fontsize=9)
    ax.set_xlabel("Tokens seen (millions)", color=INK_2)
    ax.set_ylabel("Perplexity", color=INK_2)
    fig.tight_layout()
    fig.savefig(os.path.join(run, "pretrain.png"), dpi=160, facecolor=SURFACE)


def plot_scaling(run: str):
    d = json.load(open(os.path.join(run, "scaling.json")))
    N = np.array([r["N"] for r in d["runs"]], dtype=float)
    L = np.array([r["loss"] for r in d["runs"]])
    E, A, a = d["fit"]["E"], d["fit"]["A"], d["fit"]["alpha"]
    ex = d["extrapolation"]
    fig, ax = plt.subplots(figsize=(7, 4), facecolor=SURFACE)
    _style(ax, "Scaling law — loss vs model size (D = 20·N)")
    xs = np.logspace(np.log10(N.min() * 0.8), np.log10(ex["N_main"] * 2), 200)
    ax.plot(xs, E + A * xs ** -a, color=INK_2, linewidth=1.5, linestyle="--",
            label=f"fit: L = {E:.2f} + {A:.0f}·N^-{a:.2f}")
    ax.scatter(N, L, s=48, color=SERIES["base"], edgecolor=SURFACE, linewidth=2, zorder=3, label="scaling runs")
    ax.scatter([ex["N_main"]], [ex["actual_loss"]], s=64, marker="D", color=SERIES["dpo"], edgecolor=SURFACE,
               linewidth=2, zorder=4, label="main model (actual)")
    ax.annotate(f"predicted {ex['predicted_loss']:.3f}\nactual {ex['actual_loss']:.3f}",
                (ex["N_main"], ex["actual_loss"]), xytext=(-20, 70), textcoords="offset points", ha="center",
                color=INK, fontsize=9, arrowprops=dict(arrowstyle="-", color=INK_2, lw=0.8))
    ax.set_xscale("log")
    ax.set_xlabel("Non-embedding params (N)", color=INK_2)
    ax.set_ylabel("Validation loss", color=INK_2)
    ax.legend(frameon=False, fontsize=9, labelcolor=INK_2)
    fig.tight_layout()
    fig.savefig(os.path.join(run, "scaling.png"), dpi=160, facecolor=SURFACE)


def plot_eval(run: str):
    ev = json.load(open(os.path.join(run, "eval.json")))
    m = ev["metrics"]
    models = ["base", "sft", "dpo"]
    metrics = [("keyword_hit_rate", "Keyword hit"), ("all_keywords_rate", "All 3 words"),
               ("finish_rate", "Finished (EOS)")]
    fig, (ax, ax2) = plt.subplots(1, 2, figsize=(10, 4), facecolor=SURFACE, gridspec_kw={"width_ratios": [3, 1.3]})
    _style(ax, f"Instruction following (held-out prompts, n={ev.get('n_prompts', 200)})")
    x = np.arange(len(metrics))
    w = 0.26
    for i, name in enumerate(models):
        vals = [m[name][k] * 100 for k, _ in metrics]
        bars = ax.bar(x + (i - 1) * w, vals, w * 0.92, color=SERIES[name], label=name)
        for b, v in zip(bars, vals):
            ax.text(b.get_x() + b.get_width() / 2, v + 1.5, (f"{v:.0f}" if abs(v - round(v)) < 0.05 else f"{v:.1f}"), ha="center", fontsize=8, color=INK_2)
    ax.set_xticks(x, [lbl for _, lbl in metrics])
    ax.set_ylim(0, 110)
    ax.set_ylabel("%", color=INK_2)
    ax.legend(frameon=False, fontsize=9, labelcolor=INK_2, ncol=3, loc="upper left")

    _style(ax2, "Val perplexity ↓")
    ppl = [m[n]["val_ppl"] for n in models]
    bars = ax2.bar(models, ppl, 0.6, color=[SERIES[n] for n in models])
    for b, v in zip(bars, ppl):
        ax2.text(b.get_x() + b.get_width() / 2, v + 0.08, f"{v:.2f}", ha="center", fontsize=8, color=INK_2)
    ax2.set_ylim(0, max(ppl) * 1.2)
    fig.tight_layout()
    fig.savefig(os.path.join(run, "eval.png"), dpi=160, facecolor=SURFACE)


def plot_sweep(run: str):
    path = os.path.join(run, "dpo_sweep.json")
    if not os.path.exists(path):
        return
    d = json.load(open(path))
    fig, ax = plt.subplots(figsize=(7.5, 4.4), facecolor=SURFACE)
    _style(ax, "DPO sweep — instruction following vs fluency")
    colors = {"v1": SERIES["base"], "v2": SERIES["sft"]}
    sft = d["variants"]["sft (reference)"]["metrics"]
    tol = d["selection_rule"]["fluency_ppl_max_ratio_vs_sft"]
    ax.axvline(sft["fluency_ppl"] * tol, color=INK_2, linewidth=1, linestyle=":")
    ax.annotate(f"guardrail ({tol:g}× SFT)", (sft["fluency_ppl"] * tol, 0), xytext=(4, 4),
                textcoords="offset points", fontsize=8, color=INK_2)
    ax.scatter([sft["fluency_ppl"]], [sft["all_keywords_rate"] * 100], s=60, color=INK_2, marker="s",
               edgecolor=SURFACE, linewidth=2, zorder=3)
    ax.annotate("SFT", (sft["fluency_ppl"], sft["all_keywords_rate"] * 100), xytext=(6, -12),
                textcoords="offset points", fontsize=8, color=INK_2)
    pts = [(n, v) for n, v in d["variants"].items() if not n.startswith("sft")]
    for name, v in pts:
        m = v["metrics"]
        star = name == d["best"]
        ax.scatter([m["fluency_ppl"]], [m["all_keywords_rate"] * 100], s=260 if star else 50,
                   marker="*" if star else "o", color=colors[v["reward"]], edgecolor=SURFACE,
                   linewidth=1.5 if star else 2, zorder=4)
    ax.margins(x=0.12, y=0.15)
    fig.canvas.draw()
    # Greedy label placement in pixel space: try nearby offsets, keep the first that overlaps nothing.
    cands = [(8, 6), (8, -14), (-8, 6), (-8, -14), (30, 18), (-30, 18), (30, -30), (-30, -30), (55, 0), (-55, 0)]
    boxes = []
    for _, v in pts:  # markers are obstacles too
        mx, my = ax.transData.transform((v["metrics"]["fluency_ppl"], v["metrics"]["all_keywords_rate"] * 100))
        r = 7 * fig.dpi / 72
        boxes.append((mx - r, my - r, mx + r, my + r))
    for name, v in sorted(pts, key=lambda t: -t[1]["metrics"]["all_keywords_rate"]):
        m = v["metrics"]
        label = name + (" (best)" if name == d["best"] else "")
        px, py = ax.transData.transform((m["fluency_ppl"], m["all_keywords_rate"] * 100))
        w, h = 5.2 * len(label) * fig.dpi / 100, 11 * fig.dpi / 100
        for dx, dy in cands:
            dxp, dyp = dx * fig.dpi / 72, dy * fig.dpi / 72
            x0 = px + dxp if dx > 0 else px + dxp - w
            box = (x0, py + dyp, x0 + w, py + dyp + h)
            if not any(box[0] < b[2] and b[0] < box[2] and box[1] < b[3] and b[1] < box[3] for b in boxes):
                break
        boxes.append(box)
        far = (dx, dy) not in cands[:4]
        ax.annotate(label, (m["fluency_ppl"], m["all_keywords_rate"] * 100), xytext=(dx, dy),
                    textcoords="offset points", fontsize=8, color=INK, ha="left" if dx > 0 else "right",
                    arrowprops=dict(arrowstyle="-", color=INK_2, lw=0.6) if far else None)
    for r, c in colors.items():
        ax.scatter([], [], s=40, color=c, label=f"reward {r}")
    ax.set_xlabel("Fluency ppl under base model (lower = more natural)", color=INK_2)
    ax.set_ylabel("All 3 words used (%)", color=INK_2)
    ax.legend(frameon=False, fontsize=9, labelcolor=INK_2, loc="lower right")
    fig.tight_layout()
    fig.savefig(os.path.join(run, "dpo_sweep.png"), dpi=160, facecolor=SURFACE)

if __name__ == "__main__":
    run_dir = sys.argv[1] if len(sys.argv) > 1 else "results"
    for fn, needs in ((plot_pretrain, "pretrain.json"), (plot_scaling, "scaling.json"),
                      (plot_eval, "eval.json"), (plot_sweep, "dpo_sweep.json")):
        if os.path.exists(os.path.join(run_dir, needs)):
            fn(run_dir)
        else:
            print(f"skip {fn.__name__}: {needs} not found")
    print(f"Plots saved to {run_dir}/")