#!/usr/bin/env python3
"""Render the per-config figures embedded by simulator-fidelity.md."""

from __future__ import annotations

import json
from pathlib import Path

import matplotlib

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


ROOT = Path(__file__).resolve().parents[1]
ASSET_DIR = ROOT / "docs" / "assets" / "simulator-fidelity"
DATA_PATH = ASSET_DIR / "data.json"
T0_COMPARISON_PATH = (
    ROOT / "runs" / "frontier-multicase-sufficiency-v1" / "results" / "t0-final" / "comparison.json"
)

REAL_COLOR = "#3D3D3D"
PROFILE_COLOR = "#E15759"
CALIBRATED_COLOR = "#4E79A7"


def load_data() -> dict:
    with DATA_PATH.open(encoding="utf-8") as handle:
        return json.load(handle)


def style_axis(ax: plt.Axes) -> None:
    ax.grid(axis="y", alpha=0.23, linewidth=0.8)
    ax.set_axisbelow(True)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)


def save(fig: plt.Figure, name: str) -> None:
    ASSET_DIR.mkdir(parents=True, exist_ok=True)
    fig.savefig(ASSET_DIR / f"{name}.png", dpi=220, bbox_inches="tight", facecolor="white")
    fig.savefig(ASSET_DIR / f"{name}.svg", bbox_inches="tight", facecolor="white")


def mark_top(ax: plt.Axes, x: float, y: float, text: str, color: str) -> None:
    ax.annotate(
        text,
        xy=(x, y),
        xytext=(0, 7),
        textcoords="offset points",
        ha="center",
        va="bottom",
        color=color,
        fontsize=8.5,
        weight="bold",
    )


def plot_qwen30(data: dict) -> None:
    case = data["qwen30_mixed"]
    rows = sorted(case["configs"], key=lambda row: (-row["real"], row["name"]))
    labels = [f"TP{row['tp']}\nMNS{row['mns']}" for row in rows]
    real = np.array([row["real"] for row in rows])
    profile = np.array([row["frontier_profile_only"] for row in rows])
    calibrated = np.array([row["frontier_calibrated"] for row in rows])
    x = np.arange(len(rows))
    width = 0.38

    fig, axes = plt.subplots(2, 1, figsize=(13.8, 8.7), sharex=True, sharey=True)
    panels = [
        (
            axes[0],
            profile,
            PROFILE_COLOR,
            "Frontier: operator profiles only",
            "top-set miss · Kendall τ-b = 0.000 · worst-case regret = 25.63%",
            set(case["profile_only_metrics"]["simulator_top_set"]),
        ),
        (
            axes[1],
            calibrated,
            CALIBRATED_COLOR,
            "Frontier: frozen per-TP calibration",
            "top-set hit · Kendall τ-b = 0.967 · worst-case regret = 0.76%",
            set(case["calibrated_metrics"]["simulator_top_set"]),
        ),
    ]
    real_top = set(case["calibrated_metrics"]["real_top_set"])

    for ax, simulated, color, sim_label, subtitle, simulator_top in panels:
        real_bars = ax.bar(x - width / 2, real, width, color=REAL_COLOR, label="Real vLLM")
        sim_bars = ax.bar(x + width / 2, simulated, width, color=color, label=sim_label)
        style_axis(ax)
        ax.set_ylabel("SLO-feasible req/s/GPU")
        ax.set_title(subtitle, loc="left", fontsize=11, weight="bold")
        ax.set_ylim(0, 4.35)
        ax.legend(frameon=False, ncol=2, loc="upper right")
        for idx, row in enumerate(rows):
            if row["name"] in real_top:
                real_bars[idx].set_edgecolor("#000000")
                real_bars[idx].set_linewidth(2.0)
                mark_top(ax, idx - width / 2, real[idx], "real top", REAL_COLOR)
            if row["name"] in simulator_top:
                sim_bars[idx].set_edgecolor(color)
                sim_bars[idx].set_linewidth(2.0)
                sim_bars[idx].set_hatch("///")
                mark_top(ax, idx + width / 2, simulated[idx], "sim top", color)

    axes[1].set_xticks(x, labels=labels)
    axes[1].set_xlabel("Configuration (ordered by measured real capacity)")
    fig.suptitle(
        "Qwen3-30B mixed serving: calibration changes the config-selection result",
        fontsize=14,
        weight="bold",
        y=0.995,
    )
    fig.tight_layout(rect=(0, 0, 1, 0.97))
    save(fig, "qwen30-mixed-config-ranking")
    plt.close(fig)


def plot_qwen235(data: dict) -> None:
    case = data["qwen235_prefill"]
    rows = sorted(case["configs"], key=lambda row: (-row["real"], row["name"]))
    labels = [f"TP{row['tp']} MNS{row['mns']}\nMBT{row['mbt'] // 1024}K" for row in rows]
    real = np.array([row["real"] for row in rows])
    frontier = np.array([row["frontier"] for row in rows])
    x = np.arange(len(rows))
    width = 0.38
    metrics = case["metrics"]
    real_top = set(metrics["real_top_set"])
    simulator_top = set(metrics["simulator_top_set"])

    fig, ax = plt.subplots(figsize=(13.4, 5.4))
    real_bars = ax.bar(x - width / 2, real, width, color=REAL_COLOR, label="Real community vLLM")
    sim_bars = ax.bar(x + width / 2, frontier, width, color=CALIBRATED_COLOR, label="Frontier best-effort")
    style_axis(ax)
    ax.set_ylabel("SLO-feasible req/s/GPU")
    ax.set_xticks(x, labels=labels)
    ax.set_xlabel("Configuration (ordered by measured real capacity)")
    ax.set_ylim(0, 0.088)
    ax.legend(frameon=False, ncol=2, loc="upper right")
    ax.set_title(
        "Top set matches · Spearman ρ = 0.949 · 20/20 comparable pairs agree · regret = 0",
        loc="left",
        fontsize=11,
        weight="bold",
    )
    for idx, row in enumerate(rows):
        if row["name"] in real_top:
            real_bars[idx].set_edgecolor("#000000")
            real_bars[idx].set_linewidth(2.0)
            mark_top(ax, idx - width / 2, real[idx], "real top", REAL_COLOR)
        if row["name"] in simulator_top:
            sim_bars[idx].set_edgecolor(CALIBRATED_COLOR)
            sim_bars[idx].set_linewidth(2.0)
            sim_bars[idx].set_hatch("///")
            mark_top(ax, idx + width / 2, frontier[idx], "sim top", CALIBRATED_COLOR)
    fig.suptitle(
        "Qwen3-235B-A22B-FP8 prefill-only: Frontier preserves the useful ordering",
        fontsize=14,
        weight="bold",
        y=1.01,
    )
    fig.tight_layout()
    save(fig, "qwen235-prefill-config-ranking")
    plt.close(fig)


def plot_qwen235_t0() -> None:
    with T0_COMPARISON_PATH.open(encoding="utf-8") as handle:
        comparison = json.load(handle)["comparisons"]["tpot_150ms"]
    rows = sorted(
        comparison["records"],
        key=lambda row: (row["config"]["tp"], row["config"]["mns"], row["config"]["mbt"]),
    )
    labels = [
        f"TP{row['config']['tp']} MNS{row['config']['mns']}\nMBT{row['config']['mbt'] // 1024}K"
        for row in rows
    ]
    real = np.array([row["real_capacity_per_gpu"] for row in rows])
    frontier = np.array([row["sim_capacity_per_gpu"] for row in rows])
    x = np.arange(len(rows))
    width = 0.38

    fig = plt.figure(figsize=(16.2, 5.7), constrained_layout=True)
    grid = fig.add_gridspec(1, 3, width_ratios=[2.6, 0.72, 0.72])
    ax = fig.add_subplot(grid[0, 0])
    real_bars = ax.bar(x - width / 2, real, width, color=REAL_COLOR, label="Real community vLLM")
    sim_bars = ax.bar(x + width / 2, frontier, width, color=CALIBRATED_COLOR, label="Frozen Frontier")
    style_axis(ax)
    ax.axvline(3.5, color="#999999", linewidth=1.0, linestyle="--")
    ax.set_ylabel("SLO-feasible req/s/GPU")
    ax.set_xticks(x, labels=labels)
    ax.set_ylim(0, 0.71)
    ax.legend(frameon=False, ncol=2, loc="upper right")
    ax.set_title(
        "Exact top-set match · Kendall τ-b = 0.894 · worst tie-break regret = 0",
        loc="left",
        fontsize=10.5,
        weight="bold",
    )
    real_top = set(comparison["real_top_set"])
    simulator_top = set(comparison["sim_top_set"])
    for idx, row in enumerate(rows):
        name = row["config"]["name"]
        if name in real_top:
            real_bars[idx].set_edgecolor("#000000")
            real_bars[idx].set_linewidth(1.8)
        if name in simulator_top:
            sim_bars[idx].set_edgecolor(CALIBRATED_COLOR)
            sim_bars[idx].set_linewidth(1.8)
            sim_bars[idx].set_hatch("///")

    matrices = []
    for source in ("real_capacity_per_gpu", "sim_capacity_per_gpu"):
        matrix = np.empty((2, 2))
        for row in rows:
            config = row["config"]
            if config["tp"] != 8:
                continue
            matrix[[64, 128].index(config["mns"]), [8192, 16384].index(config["mbt"])] = row[source]
        matrices.append(matrix)
    heat_axes = [fig.add_subplot(grid[0, 1]), fig.add_subplot(grid[0, 2])]
    titles = ["Real TP8", "Frontier TP8"]
    images = []
    for heat_ax, matrix, title in zip(heat_axes, matrices, titles):
        images.append(heat_ax.imshow(matrix, cmap="YlOrRd", vmin=0.15, vmax=0.30, aspect="equal"))
        heat_ax.set_title(title, fontsize=11, weight="bold")
        heat_ax.set_xticks([0, 1], labels=["8K", "16K"])
        heat_ax.set_yticks([0, 1], labels=["64", "128"])
        heat_ax.set_xlabel("MBT")
        heat_ax.set_ylabel("MNS")
        for i in range(2):
            for j in range(2):
                heat_ax.text(j, i, f"{matrix[i, j]:.2f}", ha="center", va="center", weight="bold")
    fig.colorbar(images[-1], ax=heat_axes, location="bottom", shrink=0.75, pad=0.14, label="req/s/GPU")
    fig.suptitle(
        "Qwen3-235B fixed-shape mixed: correct deployment choice, incorrect TP8 control interaction",
        fontsize=14,
        weight="bold",
    )
    save(fig, "qwen235-t0-fixed-shape-ranking")
    plt.close(fig)


def main() -> None:
    data = load_data()
    plot_qwen30(data)
    plot_qwen235(data)
    plot_qwen235_t0()


if __name__ == "__main__":
    main()