"""Hero chart for the model card: accuracy vs mean output tokens, base Dory vs Dory-Beam-4B at 1/2/4 loops.

  python3 scripts/make_chart.py   # -> release/model/assets/accuracy_vs_tokens.png

Emphasis form: Dory-Beam in the accent hue, base Dory in the de-emphasis gray. One panel per benchmark.
"""

import json
from collections import defaultdict
from pathlib import Path

import matplotlib

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

ROOT = Path(__file__).resolve().parent.parent
RES = ROOT / "data" / "eval" / "results"
OUT = ROOT / "release" / "model" / "assets" / "accuracy_vs_tokens.png"

SURFACE, INK, INK2, MUTED, GRID, AXIS = "#fcfcfb", "#0b0b0b", "#52514e", "#898781", "#e1e0d9", "#c3c2b7"
ACCENT, BASE = "#2a78d6", "#898781"
BENCHES = [("gpqa", "GPQA-Diamond"), ("math500", "MATH-500"), ("aime24", "AIME 2024"),
           ("aime25", "AIME 2025"), ("amc23", "AMC 2023"), ("heldout", "Held-out (300)")]
LOOPS = (1, 2, 4)
SERIES = [("base", "Dory1.0-4B (base)", BASE), ("mix050", "Dory-Beam-4B", ACCENT)]


def load(tag):
    d = defaultdict(list)
    for line in open(RES / f"{tag}.jsonl"):
        r = json.loads(line)
        d[r["bench"]].append(r)
    return {b: (100 * sum(x["correct"] for x in rs) / len(rs), sum(x["tokens"] for x in rs) / len(rs))
            for b, rs in d.items()}


def main():
    data = {(s, l): load(f"{s}_loops{l}-64k") for s, _, _ in SERIES for l in LOOPS}
    plt.rcParams.update({"font.family": ["Helvetica Neue", "Helvetica", "Arial", "DejaVu Sans"],
                         "font.size": 10, "axes.edgecolor": AXIS, "axes.labelcolor": INK2,
                         "xtick.color": MUTED, "ytick.color": MUTED, "text.color": INK})
    fig, axes = plt.subplots(2, 3, figsize=(12, 7.2), facecolor=SURFACE)
    fig.subplots_adjust(left=0.06, right=0.985, top=0.775, bottom=0.08, wspace=0.22, hspace=0.42)

    for ax, (bench, title) in zip(axes.flat, BENCHES):
        ax.set_facecolor(SURFACE)
        for side in ("top", "right"):
            ax.spines[side].set_visible(False)
        for side in ("left", "bottom"):
            ax.spines[side].set_linewidth(1)
        ax.grid(True, color=GRID, linewidth=1, zorder=0)
        ax.tick_params(length=0, pad=4)
        accs, toks = [], []
        for s, label, color in SERIES:
            pts = [data[(s, l)][bench] for l in LOOPS]
            ys, xs = [p[0] for p in pts], [p[1] / 1000 for p in pts]
            accs += ys
            toks += xs
            ax.plot(xs, ys, color=color, linewidth=2, solid_capstyle="round", solid_joinstyle="round", zorder=3)
            ax.plot(xs, ys, "o", color=color, markersize=8, markeredgecolor=SURFACE, markeredgewidth=2, zorder=4,
                    label=label)
            leftmost = min(range(len(xs)), key=lambda i: xs[i])
            for i, (l, x, y) in enumerate(zip(LOOPS, xs, ys)):  # loop count beside each point, in muted ink
                # the leftmost point's label goes on its left so near-coincident points don't collide
                left = i == leftmost
                ax.annotate(str(l), (x, y), xytext=(-7 if left else 7, -3), textcoords="offset points",
                            fontsize=8.5, color=MUTED, ha="right" if left else "left")
        lo, hi = min(accs), max(accs)
        pad = max(2.0, (hi - lo) * 0.18)
        ax.set_ylim(max(0, lo - pad), min(100.8, hi + pad))
        ax.set_xlim(0, max(toks) * 1.12)
        ax.xaxis.set_major_formatter(FuncFormatter(lambda v, _: "0" if v == 0 else f"{v:g}k"))
        ax.yaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v:g}%"))
        ax.set_title(title, loc="left", fontsize=11.5, fontweight="semibold", color=INK, pad=8)

    for ax in axes[1]:
        ax.set_xlabel("mean output tokens per answer", fontsize=9.5, color=INK2)
    for ax in axes[:, 0]:
        ax.set_ylabel("accuracy", fontsize=9.5, color=INK2)

    fig.text(0.06, 0.945, "Beam-distilled Dory: less than half the thinking, at a small accuracy cost",
             fontsize=17, fontweight="semibold", color=INK)
    fig.text(0.06, 0.905, "Accuracy vs mean output tokens. Numbers beside each point are the loop count (1, 2, 4). "
             "vLLM, temperature 1.0, up to 64k tokens.", fontsize=10.5, color=INK2)
    handles, labels = axes[0, 0].get_legend_handles_labels()
    fig.legend(handles, labels, loc="upper left", bbox_to_anchor=(0.053, 0.893), ncol=2, frameon=False,
               fontsize=10.5, handletextpad=0.4, columnspacing=1.6)
    OUT.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(OUT, dpi=200, facecolor=SURFACE)
    print("wrote", OUT)


if __name__ == "__main__":
    main()