#!/usr/bin/env python3
"""Plot measured knob conditional effects for the AITuner harness study."""

from __future__ import annotations

from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.lines import Line2D
from matplotlib.patches import Patch


OUT = Path("docs/harness-ablation/figures")


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


def plot_c1_surface() -> None:
    # Qwen30B mixed workload, TP x MNS screen. Values are req/s/GPU.
    # Source runs:
    # - interaction-mixed-qwen30b-tp-mns-surface-high1-dash1-d8899c5-20260701T095858Z
    # - interaction-mixed-qwen30b-tp4-mns-nocap-qps20-dash1-d8899c5-20260701T161900Z
    mns = np.array([8, 16, 32, 64])
    tp = np.array([1, 2, 4])
    values = np.array(
        [
            [2.1000, 2.3500, 2.2833, 2.2833],
            [2.2750, 2.2750, 3.2833, 3.2583],
            [1.2833, 2.4417, 2.4417, 2.4417],
        ]
    )

    fig, axes = plt.subplots(1, 2, figsize=(12.8, 4.8), gridspec_kw={"width_ratios": [1.05, 1.2]})
    ax = axes[0]
    im = ax.imshow(values, cmap="YlGnBu", aspect="auto", vmin=1.2, vmax=3.35)
    ax.set_xticks(range(len(mns)), labels=mns)
    ax.set_yticks(range(len(tp)), labels=[f"TP={x}" for x in tp])
    ax.set_xlabel("max-num-seqs (MNS)")
    ax.set_ylabel("tensor-parallel-size")
    ax.set_title("C1 response surface: req/s/GPU")
    for i in range(values.shape[0]):
        for j in range(values.shape[1]):
            color = "white" if values[i, j] > 2.75 else "black"
            ax.text(j, i, f"{values[i, j]:.2f}", ha="center", va="center", color=color, fontsize=10)
    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label("req/s/GPU")

    ax = axes[1]
    colors = {1: "#4E79A7", 2: "#59A14F", 4: "#E15759"}
    for idx, t in enumerate(tp):
        ax.plot(mns, values[idx], marker="o", linewidth=2.4, color=colors[int(t)], label=f"TP={t}")
    ax.set_xscale("log", base=2)
    ax.set_xticks(mns, labels=mns)
    ax.set_xlabel("max-num-seqs (MNS)")
    ax.set_ylabel("req/s/GPU")
    ax.set_title("Non-parallel lines imply interaction")
    ax.grid(True, axis="y", alpha=0.28)
    ax.legend(frameon=False)
    ax.annotate(
        "TP=2 only becomes best\nwhen MNS reaches 32",
        xy=(32, 3.2833),
        xytext=(20, 3.05),
        arrowprops={"arrowstyle": "->", "lw": 1.2},
        fontsize=9,
    )
    ax.annotate(
        "TP=4 is bad at MNS=8\nbut recovers at MNS>=16",
        xy=(8, 1.2833),
        xytext=(10, 1.55),
        arrowprops={"arrowstyle": "->", "lw": 1.2},
        fontsize=9,
    )
    fig.suptitle("Knob effects are conditional: MNS effect depends on TP", fontsize=14, y=1.02)
    fig.tight_layout()
    save(fig, "knob-conditional-c1-qwen30b-surface")
    plt.close(fig)


def plot_c1_oat_counterexample() -> None:
    # C1 Qwen30B: one-knob-at-a-time tuning gets trapped at a coordinate-wise
    # local optimum 25.6% below the measured global best. The right panel zooms
    # into the trap's neighbourhood so the reader can SEE that every single-knob
    # move from the trap is worse or flat, instead of having to read a caption.
    mns = [8, 16, 32, 64]
    tp = [1, 2, 4]
    values = np.array(
        [
            [2.1000, 2.3500, 2.2833, 2.2833],
            [2.2750, 2.2750, 3.2833, 3.2583],
            [1.2833, 2.4417, 2.4417, 2.4417],
        ]
    )
    idx = {(t, s): (mns.index(s), tp.index(t)) for t in tp for s in mns}

    fig, axes = plt.subplots(1, 2, figsize=(14.0, 6.2), gridspec_kw={"width_ratios": [1.5, 1.0]})
    ax = axes[0]
    im = ax.imshow(values, cmap="YlGnBu", aspect="auto", vmin=1.2, vmax=3.35)
    ax.set_xticks(range(len(mns)), labels=mns)
    ax.set_yticks(range(len(tp)), labels=[f"TP={x}" for x in tp])
    ax.set_xlabel("max-num-seqs (MNS)", fontsize=11)
    ax.set_ylabel("tensor-parallel-size", fontsize=11)
    ax.set_title("Two OAT paths from the same start", fontsize=12, loc="left")
    for i in range(values.shape[0]):
        for j in range(values.shape[1]):
            color = "white" if values[i, j] > 2.75 else "black"
            ax.text(j, i, f"{values[i, j]:.2f}", ha="center", va="center", color=color, fontsize=11, weight="bold")

    def draw_path(path: list[tuple[int, int]], color: str, labels: list[str]) -> None:
        for (a, b) in zip(path, path[1:]):
            x0, y0 = idx[a]
            x1, y1 = idx[b]
            ax.annotate(
                "",
                xy=(x1, y1),
                xytext=(x0, y0),
                arrowprops={"arrowstyle": "->", "lw": 3.2, "color": color, "shrinkA": 22, "shrinkB": 22},
            )
        for (a, b), lbl in zip(zip(path, path[1:]), labels):
            x0, y0 = idx[a]
            x1, y1 = idx[b]
            mx, my = (x0 + x1) / 2, (y0 + y1) / 2
            if x0 == x1:  # vertical move -> label to the side
                ax.text(
                    mx + 0.30,
                    my,
                    lbl,
                    color=color,
                    fontsize=9.5,
                    ha="left",
                    va="center",
                    weight="bold",
                    bbox={"boxstyle": "round,pad=0.2", "facecolor": "white", "edgecolor": "none", "alpha": 0.85},
                )
            else:  # horizontal move -> label above
                ax.text(
                    mx,
                    my - 0.32,
                    lbl,
                    color=color,
                    fontsize=9.5,
                    ha="center",
                    va="bottom",
                    weight="bold",
                    bbox={"boxstyle": "round,pad=0.2", "facecolor": "white", "edgecolor": "none", "alpha": 0.85},
                )

    # Red: tune MNS first, then TP -> walks into a coordinate-wise local optimum.
    draw_path([(1, 8), (1, 16), (4, 16)], "#C0392B", ["tune MNS", "tune TP"])
    # Green: tune TP first, then MNS -> reaches the measured global best.
    draw_path([(1, 8), (2, 8), (2, 32)], "#2E7D32", ["tune TP", "tune MNS"])

    # start / trap / best markers (explained by the legend below the grid)
    sx, sy = idx[(1, 8)]
    ax.scatter([sx], [sy], marker="o", s=210, facecolors="none", edgecolors="black", linewidths=2.2, zorder=5)
    trap = (4, 16)
    tx, ty = idx[trap]
    ax.add_patch(plt.Rectangle((tx - 0.5, ty - 0.5), 1, 1, fill=False, edgecolor="#C0392B", linewidth=3.4))
    best = (2, 32)
    bx, by = idx[best]
    ax.add_patch(plt.Rectangle((bx - 0.5, by - 0.5), 1, 1, fill=False, edgecolor="#2E7D32", linewidth=3.4))

    legend_elements = [
        Line2D([0], [0], marker="o", color="w", markerfacecolor="none", markeredgecolor="black", markeredgewidth=2, markersize=10, label="start  TP1,MNS8 = 2.10"),
        Patch(facecolor="none", edgecolor="#C0392B", linewidth=2.4, label="OAT trap  TP4,MNS16 = 2.44  (no improving single-knob move)"),
        Patch(facecolor="none", edgecolor="#2E7D32", linewidth=2.4, label="global best  TP2,MNS32 = 3.28"),
    ]
    ax.legend(handles=legend_elements, loc="upper center", bbox_to_anchor=(0.5, -0.11), ncol=1, frameon=False, fontsize=9.5)
    cbar = fig.colorbar(im, ax=ax, fraction=0.038, pad=0.035)
    cbar.set_label("req/s/GPU")

    # ---- right panel: why the red path stops (trap neighbourhood zoom) ----
    ax2 = axes[1]
    ax2.set_xlim(-0.55, 3.55)
    ax2.set_ylim(-0.65, 3.7)
    ax2.set_aspect("equal")
    ax2.set_axis_off()
    ax2.set_title("Why the red path stops here", fontsize=12, loc="left")

    # 3x3 neighbourhood of the trap (TP4,MNS16). Rows match left panel:
    # top=TP2, middle=TP4(trap), bottom=TP8(not measured). Cols: MNS 8/16/32.
    # (col, row) with row 0 at bottom.
    zoom_cells = {
        (1, 2): ("2.275", "TP2", "dead"),   # up neighbour
        (0, 1): ("1.28", "TP4", "dead"),    # left neighbour
        (1, 1): ("2.44", "TP4", "trap"),    # the trap
        (2, 1): ("2.44", "TP4", "flat"),    # right neighbour (flat, not strictly improving)
        (1, 0): ("—", "TP8", "oob"),        # down neighbour (not measured)
    }
    for (col, row), (val, tplabel, kind) in zoom_cells.items():
        x, y = col, row
        if kind == "trap":
            fc, ec, lw = "#FDECEA", "#C0392B", 3.2
        elif kind == "oob":
            fc, ec, lw = "#F2F2F2", "#CCCCCC", 1.0
        else:
            fc, ec, lw = "#FDF2F2", "#E6B0AA", 1.4
        ax2.add_patch(plt.Rectangle((x, y), 1, 1, facecolor=fc, edgecolor=ec, linewidth=lw))
        if kind == "oob":
            ax2.text(x + 0.5, y + 0.5, "no data", ha="center", va="center", fontsize=9, color="#999", style="italic")
        else:
            val_color = "#C0392B" if kind in ("dead", "flat") else "#222"
            ax2.text(x + 0.5, y + 0.62, val, ha="center", va="center", fontsize=13, weight="bold", color=val_color)
            sublabel = tplabel
            if kind == "trap":
                sublabel = f"{tplabel} · trap"
            elif kind == "flat":
                sublabel = f"{tplabel} · flat"
            ax2.text(x + 0.5, y + 0.28, sublabel, ha="center", va="center", fontsize=8.5, color="#666")
        if kind in ("dead", "flat"):
            ax2.text(x + 0.87, y + 0.87, "✗", ha="center", va="center", color="#C0392B", fontsize=18, weight="bold")

    # axis labels for the zoom
    ax2.text(-0.18, 2.5, "TP=2", ha="right", va="center", fontsize=9, color="#666")
    ax2.text(-0.18, 1.5, "TP=4", ha="right", va="center", fontsize=9, color="#666")
    ax2.text(-0.18, 0.5, "TP=8", ha="right", va="center", fontsize=9, color="#999")
    ax2.text(0.5, -0.18, "MNS=8", ha="center", va="top", fontsize=9, color="#666")
    ax2.text(1.5, -0.18, "MNS=16", ha="center", va="top", fontsize=9, color="#666")
    ax2.text(2.5, -0.18, "MNS=32", ha="center", va="top", fontsize=9, color="#666")

    ax2.text(
        1.5,
        3.45,
        "Every measured single-knob move from the trap\nis worse or flat → coordinate ascent is stuck",
        ha="center",
        va="center",
        fontsize=10.5,
        color="#222",
        bbox={"boxstyle": "round,pad=0.35", "facecolor": "white", "edgecolor": "#C0392B"},
    )

    fig.suptitle("One-knob-at-a-time tuning gets trapped: 25.6% throughput gap between two tuning orders", fontsize=14, y=1.02)
    fig.tight_layout()
    save(fig, "knob-oat-counterexample-c1-qwen30b")
    plt.close(fig)


def plot_c1_interaction_residual() -> None:
    # If TP and MNS were independent additive effects, this residual matrix would be near zero.
    mns = [8, 16, 32, 64]
    tp = [1, 2, 4]
    values = np.array(
        [
            [2.1000, 2.3500, 2.2833, 2.2833],
            [2.2750, 2.2750, 3.2833, 3.2583],
            [1.2833, 2.4417, 2.4417, 2.4417],
        ]
    )
    residual = values - values.mean(axis=1, keepdims=True) - values.mean(axis=0, keepdims=True) + values.mean()
    fig, ax = plt.subplots(figsize=(7.2, 4.8))
    limit = float(np.abs(residual).max())
    im = ax.imshow(residual, cmap="RdBu", aspect="auto", vmin=-limit, vmax=limit)
    ax.set_xticks(range(len(mns)), labels=mns)
    ax.set_yticks(range(len(tp)), labels=[f"TP={x}" for x in tp])
    ax.set_xlabel("max-num-seqs (MNS)")
    ax.set_ylabel("tensor-parallel-size")
    ax.set_title("C1 non-additive interaction residual")
    for i in range(residual.shape[0]):
        for j in range(residual.shape[1]):
            ax.text(j, i, f"{residual[i, j]:+.2f}", ha="center", va="center", fontsize=10)
    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label("req/s/GPU residual")
    fig.suptitle("Independent-knob additive model leaves large structured residuals", fontsize=13, y=1.02)
    fig.tight_layout()
    save(fig, "knob-interaction-residual-c1-qwen30b")
    plt.close(fig)


def plot_c3_lines() -> None:
    # Qwen235B decode C3, topology x MNS x MBT screen.
    # Source run:
    # interaction-qwen235b-decode-c3-topo-mns-mbt-fixed-dash1-d8899c5-20260703T022514Z
    data = {
        ("TP4/DP2/EP8", 64, 256): 0.05354166666666667,
        ("TP4/DP2/EP8", 64, 384): 0.05354166666666667,
        ("TP4/DP2/EP8", 128, 256): 0.058958333333333335,
        ("TP4/DP2/EP8", 128, 384): 0.058958333333333335,
        ("TP2/DP4/EP8", 64, 256): 0.058958333333333335,
        ("TP2/DP4/EP8", 64, 384): 0.05354166666666667,
        ("TP2/DP4/EP8", 128, 256): 0.058958333333333335,
        ("TP2/DP4/EP8", 128, 384): 0.058958333333333335,
    }
    mbt = [256, 384]
    topologies = ["TP4/DP2/EP8", "TP2/DP4/EP8"]

    fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.6), sharey=True)
    for ax, topo in zip(axes, topologies):
        for mns, color in [(64, "#4E79A7"), (128, "#F28E2B")]:
            vals = [data[(topo, mns, b)] for b in mbt]
            ax.plot(mbt, vals, marker="o", linewidth=2.6, color=color, label=f"MNS={mns}")
            for x, y in zip(mbt, vals):
                ax.text(x, y + 0.0007, f"{y:.4f}", ha="center", fontsize=9)
        ax.set_title(topo)
        ax.set_xlabel("max-num-batched-tokens (MBT)")
        ax.set_xticks(mbt)
        ax.grid(True, axis="y", alpha=0.28)
        ax.set_ylim(0.050, 0.062)
    axes[0].set_ylabel("req/s/GPU")
    axes[1].legend(frameon=False, loc="lower right")
    fig.suptitle("C3: MBT effect depends on topology and MNS", fontsize=14, y=1.02)
    fig.tight_layout()
    save(fig, "knob-conditional-c3-qwen235b-decode-lines")
    plt.close(fig)


def plot_delta_summary() -> None:
    c1_base = {
        "TP=1": (2.2833 - 2.1000) / 2.1000 * 100.0,
        "TP=2": (3.2833 - 2.2750) / 2.2750 * 100.0,
        "TP=4": (2.4417 - 1.2833) / 1.2833 * 100.0,
    }
    c3_mbt = {
        "TP4/DP2\nMNS=64": 0.0,
        "TP4/DP2\nMNS=128": 0.0,
        "TP2/DP4\nMNS=64": (0.05354166666666667 - 0.058958333333333335)
        / 0.058958333333333335
        * 100.0,
        "TP2/DP4\nMNS=128": 0.0,
    }
    c3_mns = {
        "TP4/DP2\nMBT=256": (0.058958333333333335 - 0.05354166666666667)
        / 0.05354166666666667
        * 100.0,
        "TP4/DP2\nMBT=384": (0.058958333333333335 - 0.05354166666666667)
        / 0.05354166666666667
        * 100.0,
        "TP2/DP4\nMBT=256": 0.0,
        "TP2/DP4\nMBT=384": (0.058958333333333335 - 0.05354166666666667)
        / 0.05354166666666667
        * 100.0,
    }

    panels = [
        ("C1: MNS 8->32\nunder different TP", c1_base, "#59A14F"),
        ("C3: MBT 256->384\nunder different context", c3_mbt, "#E15759"),
        ("C3: MNS 64->128\nunder different context", c3_mns, "#4E79A7"),
    ]
    fig, axes = plt.subplots(1, 3, figsize=(16, 5.2))
    for ax, (title, vals, color) in zip(axes, panels):
        labels = list(vals.keys())
        y = np.arange(len(labels))
        x = list(vals.values())
        colors = [color if v >= 0 else "#B07AA1" for v in x]
        ax.barh(y, x, color=colors)
        ax.axvline(0, color="black", linewidth=0.8)
        lo = min(x)
        hi = max(x)
        pad = max(2.0, (hi - lo) * 0.12)
        ax.set_xlim(lo - pad, hi + pad)
        ax.set_yticks(y, labels=labels, fontsize=8)
        ax.invert_yaxis()
        ax.set_xlabel("relative change in req/s/GPU (%)")
        ax.set_title(title)
        ax.grid(True, axis="x", alpha=0.25)
        for yi, xi in zip(y, x):
            ha = "left" if xi >= 0 else "right"
            offset = 0.7 if xi >= 0 else -0.7
            ax.text(xi + offset, yi, f"{xi:+.1f}%", va="center", ha=ha, fontsize=8)
    fig.suptitle("The same knob intervention has context-dependent effect size", fontsize=14, y=1.02)
    fig.tight_layout()
    save(fig, "knob-conditional-delta-summary")
    plt.close(fig)


def main() -> None:
    plot_c1_oat_counterexample()
    plot_c1_interaction_residual()
    plot_c1_surface()
    plot_c3_lines()
    plot_delta_summary()


if __name__ == "__main__":
    main()