"""Figure 10: the published root against every arm measured under the frozen protocol, on every metric."""
import csv, json, pathlib
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

S = pathlib.Path("/tmp/claude-1000/-home-mp-ubuntu-Projects-paper-qwen-intent-classifier/"
                 "0af32b5a-ac18-4e73-a3ad-1c863d0154af/scratchpad")
OUT = S / "fig2"
SURFACE = "#fcfcfb"; INK, INK2, MUTED = "#1a1a19", "#4a4a47", "#8a8a85"; GRID = "#e5e5e1"
C = {"blue": "#2a78d6", "orange": "#eb6834", "aqua": "#1baf7a", "yellow": "#eda100", "violet": "#4a3aa7"}
plt.rcParams.update({
    "figure.facecolor": SURFACE, "axes.facecolor": SURFACE, "savefig.facecolor": SURFACE,
    "axes.edgecolor": GRID, "axes.linewidth": 0.8, "axes.labelcolor": INK2,
    "xtick.color": MUTED, "ytick.color": MUTED, "text.color": INK,
    "font.size": 9, "axes.titlesize": 10.5, "axes.titleweight": "semibold",
    "grid.color": GRID, "grid.linewidth": 0.7,
    "xtick.major.size": 0, "ytick.major.size": 0,
    "legend.frameon": False, "figure.dpi": 160,
})

arms = json.loads((S / "merged_arms.json").read_text())
MI = 1024 ** 2
rows = {r["arm_id"]: r for r in csv.DictReader(open(str(S / "hf/bench/61_arm_bf16.csv")))}
g = json.loads(pathlib.Path("/home/mp_ubuntu/Projects/paper_qwen_intent_classifier/agents/sessions/"
                            "2026-09-20/S-20260920-taskblind-ladder-v1/MEASUREMENTS.json").read_text())
gridN = {"15,380": "8192", "23,551": "16384", "39,866": "32768", "72,455": "65536"}
for a in arms:
    if a["new"]:
        n = next(v for k, v in gridN.items() if k in a["arm_id"])
        a["peak_rss_mib"] = g["arms"][n]["BF16"]["peak_tree_rss_bytes"] / MI
    else:
        a["peak_rss_mib"] = float(rows[a["arm_id"]]["peak_rss_mib"])

FINAL = "39,866"
def fam(a):
    if a["new"]: return "final" if FINAL in a["arm_id"] else "grid"
    return "qwen" if a["arm_id"].startswith(("qwen35", "qwen25")) else "encoder"
COL = {"encoder": MUTED, "qwen": C["violet"], "grid": C["blue"], "final": C["orange"]}

PANELS = [
    ("macro_f1_mean", "macro F1  —  higher is better", True, False, "{:.4f}"),
    ("p50_ms", "document p50, ms  —  lower is better", False, True, "{:.1f}"),
    ("peak_rss_mib", "peak process-tree RSS, MiB  —  lower is better", False, False, "{:.0f}"),
    ("weight_file_mib", "BF16 weight file, MiB  —  lower is better", False, True, "{:.0f}"),
    ("gpu_j_per_doc", "GPU joules per document  —  lower is better", False, True, "{:.2f}"),
]

fig, axes = plt.subplots(len(PANELS), 1, figsize=(11.0, 7.4))
for ax, (key, title, higher, logx, fmt) in zip(axes, PANELS):
    ax.set_axisbelow(True)
    ax.grid(True, axis="x", alpha=0.9)
    for s in ("top", "right", "left"): ax.spines[s].set_visible(False)
    ax.spines["bottom"].set_color(GRID)
    vals = [a[key] for a in arms]
    root = next(a for a in arms if fam(a) == "final")
    order = sorted(vals, reverse=higher)
    rank = order.index(root[key]) + 1
    for a in arms:
        f = fam(a)
        if f == "final": continue
        ax.plot([a[key], a[key]], [0.18, 0.82], lw=1.6 if f == "grid" else 1.1,
                color=COL[f], alpha=0.95 if f == "grid" else 0.45, solid_capstyle="butt", zorder=3)
    ax.plot([root[key]], [0.5], marker="*", ms=17, color=C["orange"],
            markeredgecolor=SURFACE, markeredgewidth=1.2, zorder=6)
    if logx: ax.set_xscale("log")
    ax.set_yticks([]); ax.set_ylim(0, 1.25)
    ax.set_title(title, loc="left", pad=6)
    good = "top" if rank <= len(arms) / 3 else ("middle" if rank <= 2 * len(arms) / 3 else "bottom")
    ax.annotate(f"this root {fmt.format(root[key])}  ·  rank {rank} of {len(arms)}  ({good} third)",
                (0.998, 1.06), xycoords="axes fraction", ha="right", va="bottom",
                fontsize=8.5, color=C["orange"], fontweight="semibold")
    lo, hi = min(vals), max(vals)
    ax.annotate(f"best {fmt.format(order[0])}", (hi if higher else lo, 0.02),
                xycoords=("data", "axes fraction"), ha="right" if higher else "left",
                va="bottom", fontsize=7, color=MUTED)

handles = [plt.Line2D([], [], color=COL["encoder"], lw=2, label="encoder families (46)"),
           plt.Line2D([], [], color=COL["qwen"], lw=2, label="earlier Qwen arms (15)"),
           plt.Line2D([], [], color=COL["grid"], lw=2, label="task-blind vocabulary grid (4)"),
           plt.Line2D([], [], color=C["orange"], marker="*", ms=13, lw=0, label="this root: 4L, 39,866 EN/KO")]
fig.legend(handles=handles, loc="lower center", ncol=4, fontsize=8.5, bbox_to_anchor=(0.5, 0.062))
fig.suptitle("The published root against all 65 arms, on every metric the protocol measured",
             fontsize=12, fontweight="semibold", y=0.988)
fig.text(0.5, 0.030, "Each tick is one arm. Quality is a transfer probe with a fresh 14-label head, three seeds; the other four are single measured values.",
         ha="center", fontsize=7.5, color=MUTED)
fig.text(0.5, 0.008, "The four grid arms were measured in a later session; a same-host check on an earlier cut put the host effect at 1.05x on BF16 latency and 0.99x on RSS, so read small latency gaps with that in mind.",
         ha="center", fontsize=7.5, color=MUTED)
fig.subplots_adjust(left=0.035, right=0.985, top=0.925, bottom=0.135, hspace=0.72)
fig.savefig(OUT / "10_root_on_every_metric.png")
print("wrote 10_root_on_every_metric.png")