import json, glob, statistics as st
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Patch

MODELS = ["claude-opus-4-5", "claude-opus-4-6", "claude-opus-4-7", "claude-opus-4-8", "claude-opus-5"]
LABEL  = ["Opus 4.5", "Opus 4.6", "Opus 4.7", "Opus 4.8", "Opus 5"]

# --- manual verdict scoring (see write-ups) -------------------------------
# ends on the correct answer (drive) / leads with it as the first word
ENDS = {"turn1": [4/4, 3/3, 0/8, 3/3, 3/3], "turn2": [3/3, 3/3, 0/3, 2/3, 3/3]}
LEADS= {"turn1": [4/4, 3/3, 0/8, 0/3, 0/3], "turn2": [3/3, 3/3, 0/3, 1/3, 2/3]}

# --- thinking tokens straight from the run JSON ---------------------------
def think(d, m):
    fs = sorted(glob.glob(f"{d}/{m}__t*.json"))
    return st.mean([json.load(open(f))["usage"]["output_tokens_details"]["thinking_tokens"] for f in fs])
THINK = {"turn1": [think("runs", m) for m in MODELS], "turn2": [think("runs2", m) for m in MODELS]}

# --- palette --------------------------------------------------------------
INK, MUTE, GRID = "#141A1E", "#5C6B73", "#DCE4E7"
T1, T2, FAIL, PAPER = "#7FC5D2", "#0E7C93", "#C6413B", "#FFFFFF"

plt.rcParams.update({
    "font.family": "DejaVu Sans", "text.color": INK,
    "axes.edgecolor": GRID, "axes.labelcolor": MUTE,
    "xtick.color": MUTE, "ytick.color": MUTE,
    "figure.facecolor": PAPER, "axes.facecolor": PAPER,
})

fig, axes = plt.subplots(1, 3, figsize=(15.5, 5.9))
fig.subplots_adjust(left=.055, right=.985, top=.70, bottom=.16, wspace=.26)

x = range(len(MODELS)); w = 0.34

panels = [
    (axes[0], ENDS,  "Ends on the right answer",  "Says “drive” anywhere in the reply",      "pct"),
    (axes[1], LEADS, "Leads with the right answer","Says “drive” as the opening word",       "pct"),
    (axes[2], THINK, "Thinking tokens spent",     "Mean reasoning tokens before replying",   "num"),
]

for ax, data, title, sub, kind in panels:
    for i, (turn, colr) in enumerate([("turn1", T1), ("turn2", T2)]):
        vals = data[turn]
        off = (i - 0.5) * w
        bars = ax.bar([xi + off for xi in x], vals, w, color=colr,
                      edgecolor="none", zorder=3)
        top = 1.0 if kind == "pct" else max(THINK["turn2"])
        for xi, b, v in zip(x, bars, vals):
            zero = v <= 0.001
            txt = ("0%" if kind == "pct" else "0") if zero else \
                  (f"{v*100:.0f}%" if kind == "pct" else f"{v:.0f}")
            cx = b.get_x() + b.get_width()/2
            if zero:  # stub + red callout so a true zero never reads as missing data
                ax.plot([b.get_x(), b.get_x()+b.get_width()], [0, 0],
                        color=FAIL, lw=2.6, solid_capstyle="butt", zorder=4)
                ax.text(cx, top*.03, txt, ha="center", va="bottom",
                        fontsize=9, color=FAIL, fontweight="bold")
            elif kind == "pct" and v >= .98:
                pass             # sits on the top gridline; a label would only collide
            else:
                ax.text(cx, v + top*.03, txt, ha="center", va="bottom",
                        fontsize=9, color=INK)

    ax.set_title(title, fontsize=13.5, fontweight="bold", loc="left", pad=32, color=INK)
    ax.text(0, 1.018, sub, transform=ax.transAxes, fontsize=9.5, color=MUTE)
    ax.set_xticks(list(x)); ax.set_xticklabels(LABEL, fontsize=10.5, color=INK)
    ax.grid(axis="y", color=GRID, lw=.9, zorder=0)
    ax.set_axisbelow(True)
    for s in ("top", "right", "left"): ax.spines[s].set_visible(False)
    if kind == "pct":
        ax.set_ylim(0, 1.16); ax.set_yticks([0,.25,.5,.75,1])
        ax.set_yticklabels(["0", "25%", "50%", "75%", "100%"], fontsize=9.5)
    else:
        ax.set_ylim(0, max(THINK["turn2"]) * 1.22)
        ax.tick_params(axis="y", labelsize=9.5)

    # shade the 4.7 column across every panel
    ax.axvspan(2 - .5, 2 + .5, color=FAIL, alpha=.045, zorder=1)

# --- titles ---------------------------------------------------------------
fig.text(.055, .935, "The Car Wash Test", fontsize=27, fontweight="bold", color=INK)
fig.text(.055, .885,
         "“i need to get my car washed. the car wash is 50 meters away. do i walk or drive”",
         fontsize=12.5, color=MUTE, style="italic")
fig.text(.055, .845,
         "The car has to be at the car wash — the answer is drive. Turn 2 moves it 500 miles away "
         "and offers unlimited airline miles; the answer is still drive.",
         fontsize=10.5, color=MUTE)

fig.legend(handles=[Patch(facecolor=T1, label="Turn 1  ·  50 meters"),
                    Patch(facecolor=T2, label="Turn 2  ·  500 miles + free flights")],
           loc="upper right", bbox_to_anchor=(.985, .945), frameon=False,
           fontsize=10.5, ncol=1, handlelength=1.5, labelspacing=.5)

fig.text(.055, .062,
         "Opus 4.7 answers fastest, spends zero reasoning tokens, and is wrong every time — “walk” at 50 m, “fly” at 500 mi.",
         fontsize=10, color=INK)
fig.text(.055, .025,
         "n = 3 runs per model per turn, except Opus 4.7 turn 1 (n = 8) and Opus 4.5 turn 1 (n = 4).  "
         "Run through claude -p; Opus 4 and 4.1 excluded (retired).",
         fontsize=9.5, color=MUTE)

fig.savefig("The Car Wash Test.png", dpi=200, facecolor=PAPER)
print("saved")
for m, l, a, b in zip(MODELS, LABEL, THINK["turn1"], THINK["turn2"]):
    print(f"{l:<10} think t1={a:6.1f}  t2={b:6.1f}")