#!/usr/bin/env python3
"""Schematic figure frozen before EXP-TP2-PREFILL-SERVING GPU execution."""

from pathlib import Path

import matplotlib

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

ROOT = Path(__file__).resolve().parent
fig, axes = plt.subplots(1, 2, figsize=(9.6, 4.0), dpi=160)
colors = {"real": "#111111", "sim": "#d95f02", "candidate": "#1b75bc"}

labels = ["TP2", "TP4"]
x = np.arange(2)
width = 0.34
axes[0].bar(x - width / 2, [410, 231], width, color=colors["real"], label="real anchor")
axes[0].bar(x + width / 2, [356.25, 231.33], width, color=colors["sim"], label="structured sim")
axes[0].set_xticks(x, labels)
axes[0].set_ylabel("q8k / ctx0 chunk time (ms)")
axes[0].set_title("(a) Entry residual is TP2-only", loc="left", fontsize=10)
axes[0].legend(frameon=False, fontsize=8)

arms = ["A\nstandalone", "B\nwarm+route", "C\nserving"]
expected = [171.1, 205, 225]
axes[1].bar(np.arange(3), expected, color=[colors["sim"], colors["candidate"], colors["real"]])
axes[1].axhline(224.9, color="#777777", linestyle="--", linewidth=1, label="needed to close 54 ms")
axes[1].set_xticks(np.arange(3), arms)
axes[1].set_ylabel("TP2 MoE per chunk (ms)")
axes[1].set_title("(b) Three-arm decision (schematic)", loc="left", fontsize=10)
axes[1].legend(frameon=False, fontsize=8)

for ax in axes:
    ax.spines[["top", "right"]].set_visible(False)
    ax.grid(axis="y", color="#dedbd2", linewidth=0.8)
    ax.set_axisbelow(True)

fig.suptitle(
    "MOCK / schematic — EXP-TP2-PREFILL-SERVING (B/C not measured)",
    x=0.01,
    ha="left",
    color=colors["sim"],
    fontsize=9,
)
fig.tight_layout(rect=(0, 0, 1, 0.95))
fig.savefig(ROOT / "figure-prototype.png")
fig.savefig(ROOT / "figure-prototype.svg")
print(ROOT / "figure-prototype.png")
