"""Comparative figures for milestone-gqa-decode-long-ctx-4cases.

Reads sweep.json (emitted by ``kernbench run --bench
milestone-gqa-decode-long-ctx-4cases``) and writes three PNGs into
``docs/report/1H-codesign-paper/figures/``:

  gqa_decode_long_ctx_4cases_latency.png   end-to-end latency per case
  gqa_decode_long_ctx_4cases_traffic.png   ipcq/dma op-count breakdown
  gqa_decode_long_ctx_4cases_memory.png    KV bytes per cube per case

Run (after the bench):
    GQA_DECODE_LONG_CTX_4CASES_RUN=1 python -m kernbench.cli.main run \\
        --bench milestone-gqa-decode-long-ctx-4cases --topology topology.yaml
    python scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.py
"""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

_REPO_ROOT = Path(__file__).resolve().parents[2]
_FIG_DIR = _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
_SWEEP_JSON = (
    _REPO_ROOT / "src" / "kernbench" / "benches"
    / "1H_milestone_output" / "gqa_decode_long_ctx_4cases" / "sweep.json"
)

# Panel name → (short label, case ordinal for left-to-right plot order).
_CASE_INFO = {
    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp": (
        "Case 1\nCube-SP × PE-TP", 1),
    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp": (
        "Case 2\nCube-Repl × PE-TP", 2),
    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp": (
        "Case 3\nCube-Repl × PE-SP", 3),
    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp": (
        "Case 4 ★\nCube-SP × PE-SP", 4),
}


def _load() -> list[dict]:
    return json.loads(_SWEEP_JSON.read_text())["rows"]


def _sorted_by_case(rows: list[dict]) -> list[dict]:
    return sorted(rows, key=lambda r: _CASE_INFO[r["panel"]][1])


def _plot_latency(rows: list[dict]) -> Path:
    rows = _sorted_by_case(rows)
    labels = [_CASE_INFO[r["panel"]][0] for r in rows]
    lat_us = [r["latency_ns"] / 1e3 for r in rows]
    colors = ["#888", "#888", "#888", "#3b6ea5"]  # Case 4 highlighted
    fig, ax = plt.subplots(figsize=(8.0, 4.5))
    bars = ax.bar(labels, lat_us, color=colors, width=0.6)
    ax.set_ylabel("end-to-end latency (µs)")
    ax.set_title(
        "Long-context decode 4-cases — end-to-end latency per case\n"
        "LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs)"
    )
    ax.bar_label(bars, fmt="%.1f", padding=3, fontsize=9)
    ax.grid(axis="y", ls=":", alpha=0.5)
    ax.set_ylim(0, max(lat_us) * 1.15)
    fig.tight_layout()
    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_latency.png"
    fig.savefig(out, dpi=150)
    plt.close(fig)
    return out


def _plot_traffic(rows: list[dict]) -> Path:
    rows = _sorted_by_case(rows)
    labels = [_CASE_INFO[r["panel"]][0] for r in rows]
    x = list(range(len(rows)))
    keys = ["ipcq_copy_count", "dma_read_count", "dma_write_count"]
    disp = ["IPCQ copy", "DMA read", "DMA write"]
    colors = ["#c0504d", "#9bbb59", "#8064a2"]
    w = 0.25
    fig, ax = plt.subplots(figsize=(9.0, 4.5))
    for i, (k, d, c) in enumerate(zip(keys, disp, colors)):
        vals = [r["op_log_summary"][k] for r in rows]
        ax.bar([xi + (i - 1) * w for xi in x], vals, width=w, label=d, color=c)
    ax.set_xticks(list(x))
    ax.set_xticklabels(labels, fontsize=9)
    ax.set_ylabel("op count")
    ax.set_title("Long-context decode 4-cases — op-count breakdown per case")
    ax.legend(fontsize=9)
    ax.grid(axis="y", ls=":", alpha=0.5)
    fig.tight_layout()
    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_traffic.png"
    fig.savefig(out, dpi=150)
    plt.close(fig)
    return out


def _kv_bytes_per_cube(panel: str, *, S_kv: int, h_kv: int,
                       d_head: int, C: int) -> int:
    """KV bytes a single cube's HBM holds (K + V together, f16)."""
    # Cube-replicate ⇒ each cube holds full S_kv.
    # Cube-SP        ⇒ each cube holds S_kv / C.
    # PE replicate vs row_wise share the cube's HBM and don't change
    # per-cube bytes.
    S_per_cube = S_kv if "cube_repl" in panel else S_kv // C
    return 2 * S_per_cube * h_kv * d_head * 2  # K + V, f16 (2 B/elem)


def _plot_memory(rows: list[dict]) -> Path:
    rows = _sorted_by_case(rows)
    labels = [_CASE_INFO[r["panel"]][0] for r in rows]
    mib_per_cube = [
        _kv_bytes_per_cube(
            r["panel"], S_kv=r["S_kv"], h_kv=r["h_kv"],
            d_head=r["d_head"], C=r["C"],
        ) / (1024 * 1024)
        for r in rows
    ]
    # SP = blue (efficient); Repl = red (wasteful).
    colors = ["#3b6ea5", "#c0504d", "#c0504d", "#3b6ea5"]
    fig, ax = plt.subplots(figsize=(8.0, 4.5))
    bars = ax.bar(labels, mib_per_cube, color=colors, width=0.6)
    ax.set_ylabel("KV bytes per cube (MiB, K + V, f16)")
    ax.set_title(
        "Long-context decode 4-cases — KV memory per cube\n"
        "(one KV-head group; per-layer, per-token state)"
    )
    ax.bar_label(bars, fmt="%.1f", padding=3, fontsize=9)
    ax.grid(axis="y", ls=":", alpha=0.5)
    ax.set_ylim(0, max(mib_per_cube) * 1.15)
    fig.tight_layout()
    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_memory.png"
    fig.savefig(out, dpi=150)
    plt.close(fig)
    return out


def main() -> None:
    rows = _load()
    _FIG_DIR.mkdir(parents=True, exist_ok=True)
    p1 = _plot_latency(rows)
    p2 = _plot_traffic(rows)
    p3 = _plot_memory(rows)
    print(f"wrote {p1}")
    print(f"wrote {p2}")
    print(f"wrote {p3}")


if __name__ == "__main__":
    main()