#!/usr/bin/env python3
"""Render the profile ablation and execution-context diagnostics."""

from __future__ import annotations

import argparse
import json
from collections import Counter
from pathlib import Path

import matplotlib

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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--s2", type=Path, required=True)
    parser.add_argument("--routing", type=Path, required=True)
    parser.add_argument("--opprof", type=Path, required=True)
    parser.add_argument("--p1", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    s2 = json.loads(args.s2.read_text())
    routing = json.loads(args.routing.read_text())
    opprof = json.loads(args.opprof.read_text())
    p1 = json.loads(args.p1.read_text())

    plt.rcParams.update({"font.size": 9, "axes.titlesize": 10, "axes.labelsize": 9})
    fig, axes = plt.subplots(2, 2, figsize=(13.2, 8.2), constrained_layout=True)

    real = s2["real_scores"]
    cells = sorted(real, key=lambda cell: (-real[cell], cell))
    x = np.arange(len(cells))
    series = (
        ("Real", real, "#111111", "o"),
        (
            "Old profile-only",
            s2["historical_modes"]["historical-profile-only"]["simulated_scores"],
            "#d95f02",
            "s",
        ),
        (
            "vLLM 0.20 profile-only",
            s2["vllm020_profile_only"]["SLO-gated"]["simulated_scores"],
            "#7570b3",
            "^",
        ),
        (
            "Frozen per-TP calibration",
            s2["historical_modes"]["historical-per-tp-calibration"]["simulated_scores"],
            "#1b9e77",
            "D",
        ),
    )
    ax = axes[0, 0]
    for label, values, color, marker in series:
        ax.plot(x, [values[cell] for cell in cells], label=label, color=color, marker=marker, lw=1.7)
    ax.set_xticks(x, [cell.replace("_", "\n") for cell in cells])
    ax.set_ylabel("SLO-feasible req/s/GPU")
    ax.set_title("(a) Full 92-probe config ranking")
    ax.grid(axis="y", alpha=0.25)
    ax.legend(fontsize=8, ncol=2, loc="upper right")

    ax = axes[0, 1]
    categories = ("Actual prefill", "Actual decode", "Frontier prior")
    cv = [
        routing["phase_summary"]["prefill"]["actual"]["load_cv"]["median"],
        routing["phase_summary"]["decode"]["actual"]["load_cv"]["median"],
        routing["phase_summary"]["prefill"]["frontier_simulation"]["load_cv"]["median"],
    ]
    max_ratio = [
        routing["phase_summary"]["prefill"]["actual"]["max_load_ratio"]["median"],
        routing["phase_summary"]["decode"]["actual"]["max_load_ratio"]["median"],
        routing["phase_summary"]["prefill"]["frontier_simulation"]["max_load_ratio"]["median"],
    ]
    bx = np.arange(len(categories))
    width = 0.34
    bars1 = ax.bar(bx - width / 2, cv, width, label="Load CV", color="#66c2a5")
    bars2 = ax.bar(bx + width / 2, max_ratio, width, label="Max/mean load", color="#fc8d62")
    ax.bar_label(bars1, fmt="%.2f", fontsize=8)
    ax.bar_label(bars2, fmt="%.2f", fontsize=8)
    ax.set_xticks(bx, categories)
    ax.set_ylim(0, max(max_ratio) * 1.2)
    ax.set_title("(b) Per-layer MoE routing skew")
    ax.legend(fontsize=8)
    ax.grid(axis="y", alpha=0.25)

    graph_counts = {phase: Counter() for phase in ("pure_decode", "pure_prefill", "true_mixed")}
    for cell in opprof["cells"]:
        for group in cell["groups"]:
            graph_counts[group["phase"]][group["cudagraph_runtime_mode"]] += int(group["steps"])
    ax = axes[1, 0]
    phases = tuple(graph_counts)
    bottoms = np.zeros(len(phases))
    for mode, color in (("FULL", "#1b9e77"), ("PIECEWISE", "#e6ab02"), ("NONE", "#d95f02")):
        values = np.asarray(
            [100 * graph_counts[phase][mode] / sum(graph_counts[phase].values()) for phase in phases]
        )
        ax.bar(np.arange(len(phases)), values, bottom=bottoms, label=mode, color=color)
        bottoms += values
    ax.set_xticks(np.arange(len(phases)), [phase.replace("_", " ") for phase in phases])
    ax.set_ylabel("Observed scheduler steps (%)")
    ax.set_ylim(0, 100)
    ax.set_title("(c) Real vLLM execution mode is phase-dependent")
    ax.legend(fontsize=8, ncol=3, loc="lower left")

    ax = axes[1, 1]
    mode_order = ("historical-calibrated", "historical-profile-only", "vllm020-profile-only")
    labels = ("Per-TP\ncalibration", "Old\nprofile-only", "vLLM 0.20\nprofile-only")
    accuracy = [100 * p1["summaries"][mode]["probe_classification"]["accuracy"] for mode in mode_order]
    mae = [100 * p1["summaries"][mode]["pass_rate_mae"] for mode in mode_order]
    px = np.arange(len(labels))
    bars1 = ax.bar(px - width / 2, accuracy, width, label="Label accuracy (higher better)", color="#1b9e77")
    bars2 = ax.bar(px + width / 2, mae, width, label="Pass-rate MAE (lower better)", color="#d95f02")
    ax.bar_label(bars1, fmt="%.1f", fontsize=8)
    ax.bar_label(bars2, fmt="%.1f", fontsize=8)
    ax.set_xticks(px, labels)
    ax.set_ylabel("Percent")
    ax.set_ylim(0, 105)
    ax.set_title("(d) Held-out P1 boundary probes (12 labels)")
    ax.legend(fontsize=8, loc="upper right")
    ax.grid(axis="y", alpha=0.25)

    fig.suptitle(
        "Qwen3-30B-A3B / vLLM 0.20 / BF16 / dash0 H20: operator provenance is not execution-context fidelity",
        fontsize=12,
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(args.output, dpi=180)
    fig.savefig(args.output.with_suffix(".svg"))
    print(args.output)


if __name__ == "__main__":
    main()