"""Bar plot for the multi-model GQA composite bench.

Topology per model: cubes per KV group C = h_q, P = 8 PEs / cube.

Reads ``sweep_decode_models.json`` (produced by
``gqa_decode_long_ctx_models.py``) and writes a two-panel PNG
comparing end-to-end latency and PE_CPU dispatch across six GQA models
where each model uses a topology sized to match its G value
(cubes per KV group = h_q per KV group; PEs per cube = 8):

  Left panel  — end-to-end decode latency (µs), one bar per model, sorted
                by G. Bar labels show C, N, and latency.
  Right panel — PE_CPU command count per model.

Run (after the bench):
    python scripts/paper/paper_plot_gqa_decode_models.py
"""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib

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

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

# Model display labels
_MODEL_LABELS = {
    "gemma2-27b":       "Gemma 2 27B\n(G=2)",
    "llama3-8b":        "LLaMA-3 8B\n(G=4)",
    "qwen2.5-7b":       "Qwen 2.5 7B\n(G=7)",
    "llama3-70b":       "LLaMA-3 70B\n(G=8)",
    "qwen2.5-72b":      "Qwen 2.5 72B\n(G=8)",
    "command-r-plus":   "Command R+\n(G=12)",
}

# Bar color per model family
_FAMILY_COLOUR = {
    "Google":  "#4f8a4f",
    "Meta":    "#3b6ea5",
    "Alibaba": "#c86432",
    "Cohere":  "#8b5a2b",
}


def main() -> None:
    sweep = json.loads(_SWEEP_JSON.read_text())
    rows_by_model = {r["model"]: r for r in sweep["rows"]}
    # Sort models by G (ascending) so the trend is left-to-right.
    ordered = sorted(_MODEL_LABELS.keys(), key=lambda m: rows_by_model[m]["G"])

    fig, (ax_lat, ax_cmd, ax_brk) = plt.subplots(1, 3, figsize=(18.5, 5.0))
    xs = np.arange(len(ordered))

    # ── Left: latency ─────────────────────────────────────────────
    latencies = [rows_by_model[m]["latency_ns"] / 1e3 for m in ordered]
    colours = [_FAMILY_COLOUR[rows_by_model[m]["family"]] for m in ordered]
    bars = ax_lat.bar(xs, latencies, 0.62,
                      color=colours, edgecolor="#222", linewidth=0.7)
    for i, m in enumerate(ordered):
        r = rows_by_model[m]
        # Latency label above bar
        ax_lat.text(i, latencies[i] + max(latencies) * 0.015,
                    f"{latencies[i]:.1f} µs",
                    ha="center", va="bottom", fontsize=9,
                    fontweight="bold", color="#222")
        # Topology label inside bar
        ax_lat.text(i, latencies[i] / 2,
                    f"cubes = {r['C']}",
                    ha="center", va="center", fontsize=9, color="white")

    ax_lat.set_xticks(xs)
    ax_lat.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9)
    ax_lat.set_ylabel("end-to-end decode latency (µs)")
    ax_lat.set_title(
        f"Composite Case-6 decode latency at $S_{{kv}}=128$K\n"
        "topology per model: cubes per KV group = h_q of the KV group",
        fontsize=11,
    )
    ax_lat.grid(True, axis="y", ls=":", alpha=0.5)
    ax_lat.set_ylim(0, max(latencies) * 1.18)

    # Legend (family)
    from matplotlib.patches import Patch
    seen_families = []
    handles = []
    for m in ordered:
        fam = rows_by_model[m]["family"]
        if fam not in seen_families:
            seen_families.append(fam)
            handles.append(Patch(facecolor=_FAMILY_COLOUR[fam],
                                 edgecolor="#222", label=fam))
    ax_lat.legend(handles=handles, loc="upper left", fontsize=9,
                  title="family")

    # ── Right: PE_CPU dispatch ─────────────────────────────────────
    cmds = [rows_by_model[m]["pe_cpu_cmd_count"] for m in ordered]
    bars2 = ax_cmd.bar(xs, cmds, 0.62,
                       color=colours, edgecolor="#222", linewidth=0.7)
    for i, m in enumerate(ordered):
        ax_cmd.text(i, cmds[i] + max(cmds) * 0.02,
                    f"{cmds[i]}",
                    ha="center", va="bottom", fontsize=9,
                    fontweight="bold", color="#222")
        r = rows_by_model[m]
        ax_cmd.text(i, cmds[i] / 2,
                    f"cubes = {r['C']}",
                    ha="center", va="center", fontsize=9, color="white")

    ax_cmd.set_xticks(xs)
    ax_cmd.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9)
    ax_cmd.set_ylabel("PE_CPU commands issued")
    ax_cmd.set_title(
        "PE_CPU dispatch (composite, per KV group)\n"
        "growing sub-mesh width or height → more reduce hops",
        fontsize=11,
    )
    ax_cmd.grid(True, axis="y", ls=":", alpha=0.5)
    ax_cmd.set_ylim(0, max(cmds) * 1.18)

    # ── Third: matmul vs comm (op-kind occupancy) ─────────────────────
    # Only the two kinds that carry useful work — matmul (GEMM + MATH
    # occupancy summed across all engines) and comm (DMA occupancy).
    # These are op-log sums across components, not critical-path
    # attribution; parallelism means they don't sum to wall-clock
    # latency. Bars are absolute µs so the reduce-cost growth with C
    # is visible.
    matmul_us = [rows_by_model[m].get("matmul_ns", 0.0) / 1e3 for m in ordered]
    comm_us = [rows_by_model[m].get("comm_ns", 0.0) / 1e3 for m in ordered]
    max_v = max(max(matmul_us), max(comm_us))
    width = 0.35

    ax_brk.bar(xs - width / 2, matmul_us, width,
               color="#3b6ea5", edgecolor="#222", linewidth=0.6,
               label="matmul (GEMM + MATH)")
    ax_brk.bar(xs + width / 2, comm_us, width,
               color="#c0504d", edgecolor="#222", linewidth=0.6,
               label="communication (DMA)")
    for i in range(len(ordered)):
        ax_brk.text(xs[i] - width / 2, matmul_us[i] + max_v * 0.015,
                    f"{matmul_us[i]:.1f}",
                    ha="center", va="bottom", fontsize=8,
                    color="#222", fontweight="bold")
        ax_brk.text(xs[i] + width / 2, comm_us[i] + max_v * 0.015,
                    f"{comm_us[i]:.1f}",
                    ha="center", va="bottom", fontsize=8,
                    color="#222", fontweight="bold")

    ax_brk.set_xticks(xs)
    ax_brk.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9)
    ax_brk.set_ylabel("op-log occupancy (µs, summed across engines)")
    ax_brk.set_title(
        "Matmul vs communication engine work per model\n"
        "sum of GEMM/MATH and DMA occupancy — not wall-clock (overlap)",
        fontsize=11,
    )
    ax_brk.grid(True, axis="y", ls=":", alpha=0.5)
    ax_brk.set_ylim(0, max_v * 1.20)
    ax_brk.legend(loc="upper left", fontsize=9)

    fig.suptitle(
        "Multi-model attention comparison — Case-6 composite kernel"
        "  ($S_{kv}=128$K, $T_q=1$, P=8 PEs / cube)",
        fontsize=12, fontweight="bold",
    )
    fig.tight_layout(rect=(0, 0, 1, 0.93))

    out = _FIG_DIR / "gqa_decode_models.png"
    fig.savefig(out, dpi=150)
    plt.close(fig)
    print(f"wrote {out}")

    if _PAPER_FIG_DIR.is_dir():
        dst = _PAPER_FIG_DIR / out.name
        dst.write_bytes(out.read_bytes())
        print(f"copied {dst}")


if __name__ == "__main__":
    main()