"""milestone-gqa-llama70b bench: GQA Llama-70B 4-panel sweep (ADR-0057 v1).

Self-contained eval bench (ADR-0054). Drives the four panels of the GQA
Llama-70B sharding study through ``run_bench`` with ``enable_data=True``,
harvests op_log summaries, and writes JSON into
``benches/1H_milestone_output/gqa/sweep.json``.

v1 (sub-cycle 4a + 4c.0) covers all four panels at validation scale:

  Panel name in JSON / test    Study label   SFR install used
  ─────────────────────────────────────────────────────────────────────
  single_user_prefill          TL            configure_sfr_intracube_pe_ring
  multi_user_prefill           TR            configure_sfr_intercube_multisip
  single_user_decode           BL            configure_sfr_intracube_pe_ring
  multi_user_decode            BR            configure_sfr_intercube_multisip

Per the GQA sharding study (`llm_paper_review/notes/GQA_MHA_sharding/scripts
/_gen_llama70b_1M_4cases.py`):

  Single User (B=1) panels — TL prefill, BL decode:
    "n_cubes: 8 (1 KV-group)", KV split @ PEs intra-cube. Each cube does
    its own 8-PE ring with no cube-to-cube attention traffic. At Llama-70B
    headline scale this is 64 cubes (8 KV-groups × 8 cubes/group), each
    independently running the kernel for one Q-head; 1 user spans all 64.

  Multi User (B=8) panels — TR prefill, BR decode:
    "8 cubes / KV-group", KV split @ cubes inter-cube. The 8 cubes of a
    KV-group form a ring; "Inside each cube: 8 PEs each handle 1 different
    user → Batch on batch, batch = 8/cube." At headline scale 8 KV-groups
    run in parallel = 64 cubes serving 8 users.

Kernels use the mesh-native variants (ADR-0059), invoked with the
``rank_axis`` kwarg (0 for single_user PE-level rings, 1 for multi_user
cube-level rings). The v1 multi_user kernel gates ``pe_id != 0`` to return,
which simplifies B=8 → B=1 — that's a validation simplification; the
per-cube "Batch on batch" parallelism is sub-cycle 4c headline work.

Validation-scale config (ADR-0057 D4) — kept small so the simulator's
1 MB per-PE TCM scratch budget is not exhausted across n_ranks ring steps:
``S_q_prefill = S_kv_per_rank = 16``, ``h_q = h_kv = 1``, ``d_head = 64``,
``n_ranks_single_user = 8`` (8 PEs of one cube), ``n_ranks_multi_user = 4``
(half a KV-group, vs the study's 8). Headline-scale dims (``S_q = 1M``,
``S_kv = 1M``, ``h_q = 8 / h_kv = 1`` GQA, ``d_head = 128``, ``B = 8``) are
also deferred.
"""
from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

from kernbench.benches._attention_mesh_kv import attention_mesh_kv_kernel
from kernbench.benches._attention_mesh_mlo import attention_mesh_mlo_kernel
from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_kernel
from kernbench.benches.registry import bench
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import (
    configure_sfr_intercube_multisip,
    configure_sfr_intracube_pe_ring,
)
from kernbench.policy.placement.dp import DPPolicy

_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa"
_SWEEP_JSON = _OUTPUT_DIR / "sweep.json"

# ── Validation-scale config (ADR-0057 D4) ─────────────────────────────

_S_Q_PREFILL = 16
_S_Q_DECODE = 1
_S_KV_PER_RANK = 16
_H_Q = 1
_H_KV = 1
_D_HEAD = 64
_N_RANKS_SINGLE_USER = 8
_N_RANKS_MULTI_USER = 4
_DTYPE = "f16"

_PANELS_V1 = (
    "single_user_prefill",
    "multi_user_prefill",
    "single_user_decode",
    "multi_user_decode",
)

# Panel → (kernel, SFR install, S_q, n_ranks, rank_axis, mesh_shape)
# ``mesh_shape`` is ``None`` for 1D-ring kernels and ``(rows, cols)`` for the
# 2D row-then-col AllReduce-mlo kernel (multi_user_decode); when set, the
# launch passes ``(mesh_rows, mesh_cols)`` instead of ``n_ranks``.
_PANEL_DISPATCH: dict[
    str, tuple[Any, Any, int, int, int, tuple[int, int] | None]
] = {
    "single_user_prefill": (
        attention_mesh_kv_kernel, configure_sfr_intracube_pe_ring,
        _S_Q_PREFILL, _N_RANKS_SINGLE_USER, 0, None,
    ),
    "multi_user_prefill": (
        attention_mesh_kv_kernel, configure_sfr_intercube_multisip,
        _S_Q_PREFILL, _N_RANKS_MULTI_USER, 1, None,
    ),
    "single_user_decode": (
        attention_mesh_mlo_kernel, configure_sfr_intracube_pe_ring,
        _S_Q_DECODE, _N_RANKS_SINGLE_USER, 0, None,
    ),
    # multi_user_decode uses the C2 2D AllReduce-mlo kernel. (1, 4)
    # degenerates to a row-only AllReduce equivalent to the prior 1D ring
    # at n_ranks=4 — no op_log_summary regression. Headline 8-cube
    # KV-groups land at (2, 4).
    "multi_user_decode": (
        attention_mesh_mlo_2d_kernel, configure_sfr_intercube_multisip,
        _S_Q_DECODE, _N_RANKS_MULTI_USER, 1, (1, _N_RANKS_MULTI_USER),
    ),
}


# ── Per-panel bench fn ─────────────────────────────────────────────────


def _make_bench_fn(panel: str):
    kernel, sfr_install, S_q, n_ranks, rank_axis, mesh_shape = (
        _PANEL_DISPATCH[panel]
    )
    is_multi_user = panel.startswith("multi_user_")

    def _bench_fn(ctx):
        sfr_install(
            ctx.engine, ctx.spec,
            resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"),
        )
        if is_multi_user:
            dp_full = DPPolicy(
                cube="replicate", pe="replicate",
                num_cubes=n_ranks, num_pes=8,
            )
            dp_kv = DPPolicy(
                cube="row_wise", pe="replicate",
                num_cubes=n_ranks, num_pes=8,
            )
        else:
            dp_full = DPPolicy(
                cube="replicate", pe="replicate",
                num_cubes=1, num_pes=n_ranks,
            )
            dp_kv = DPPolicy(
                cube="replicate", pe="row_wise",
                num_cubes=1, num_pes=n_ranks,
            )
        q = ctx.zeros((S_q, _H_Q * _D_HEAD),
                      dtype=_DTYPE, dp=dp_full, name=f"{panel}_q")
        k = ctx.zeros((_S_KV_PER_RANK * n_ranks, _H_KV * _D_HEAD),
                      dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
        v = ctx.zeros((_S_KV_PER_RANK * n_ranks, _H_KV * _D_HEAD),
                      dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
        o = ctx.empty((S_q, _H_Q * _D_HEAD),
                      dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
        # rank_axis is a positional arg; _auto_dim_remap=False keeps
        # d_head=64 from colliding with the multi_user K's global M=64.
        if mesh_shape is None:
            ctx.launch(
                f"{panel}_mesh", kernel,
                q, k, v, o,
                S_q, _S_KV_PER_RANK, _H_Q, _H_KV, _D_HEAD, n_ranks,
                rank_axis,
                _auto_dim_remap=False,
            )
        else:
            mesh_rows, mesh_cols = mesh_shape
            ctx.launch(
                f"{panel}_mesh", kernel,
                q, k, v, o,
                S_q, _S_KV_PER_RANK, _H_Q, _H_KV, _D_HEAD,
                mesh_rows, mesh_cols,
                rank_axis,
                0,  # cube_start=0: this panel's launch starts at cube 0
                _auto_dim_remap=False,
            )

    return _bench_fn


# ── Op-log summary harvest ─────────────────────────────────────────────


def _summarize_op_log(op_log) -> dict[str, int]:
    """Counts per ADR-0057 D7 op_log_summary contract."""
    gemm_count = 0
    ipcq_send_count = 0
    ipcq_recv_count = 0
    dma_read_count = 0
    dma_write_count = 0
    for r in op_log:
        if r.op_kind == "gemm":
            gemm_count += 1
        elif r.op_name == "dma_read":
            dma_read_count += 1
        elif r.op_name == "dma_write":
            dma_write_count += 1
        elif r.op_name == "ipcq_send":
            ipcq_send_count += 1
        elif r.op_name == "ipcq_recv":
            ipcq_recv_count += 1
        elif r.op_name == "ipcq_copy":
            # The inbound DMA records ipcq_copy (one per send/recv pair).
            # Count it as both a send and a recv side so the row's
            # ipcq_send_count and ipcq_recv_count are non-zero even when
            # the engine logs the collective via the inbound copy alone.
            ipcq_send_count += 1
            ipcq_recv_count += 1
    return {
        "gemm_count": gemm_count,
        "ipcq_send_count": ipcq_send_count,
        "ipcq_recv_count": ipcq_recv_count,
        "dma_read_count": dma_read_count,
        "dma_write_count": dma_write_count,
    }


def _run_panel(panel: str, topology: str) -> dict:
    """Run one panel via a fresh engine; return its row dict."""
    from kernbench.runtime_api.bench_runner import run_bench
    from kernbench.runtime_api.types import resolve_device
    from kernbench.sim_engine.engine import GraphEngine
    from kernbench.topology.builder import resolve_topology

    topo = resolve_topology(topology)
    result = run_bench(
        topology=topo, bench_fn=_make_bench_fn(panel),
        device=resolve_device(None),
        engine_factory=lambda t, d: GraphEngine(
            getattr(t, "topology_obj", t), enable_data=True,
        ),
    )
    if not result.completion.ok:
        raise RuntimeError(
            f"milestone-gqa-llama70b panel {panel!r} failed: {result.completion}"
        )
    _, _, _, n_ranks, _, _ = _PANEL_DISPATCH[panel]
    return {
        "panel": panel,
        "n_ranks": n_ranks,
        "op_log_summary": _summarize_op_log(result.engine.op_log),
    }


# ── Figure renderers (sub-cycle 4c, 5 of 6 figures) ──────────────────
#
# Sixth figure ``gqa_scaling.png`` is deferred to after sub-cycle 4b
# lands the Q/cube ∈ {1, 2, 4} sweep on multi_user_* panels — it needs
# multiple sweep.json rows per multi_user panel to be meaningful.

_OP_LOG_KEYS = (
    "gemm_count",
    "ipcq_send_count",
    "ipcq_recv_count",
    "dma_read_count",
    "dma_write_count",
)
_OP_LOG_DISPLAY = {
    "gemm_count":      "GEMM",
    "ipcq_send_count": "IPCQ send",
    "ipcq_recv_count": "IPCQ recv",
    "dma_read_count":  "DMA read",
    "dma_write_count": "DMA write",
}
_OP_LOG_COLORS = {
    "gemm_count":      "#F59E0B",
    "ipcq_send_count": "#3B82F6",
    "ipcq_recv_count": "#10B981",
    "dma_read_count":  "#A855F7",
    "dma_write_count": "#EF4444",
}
_PANEL_DISPLAY = {
    "single_user_prefill": "single_user / prefill",
    "multi_user_prefill":  "multi_user / prefill",
    "single_user_decode":  "single_user / decode",
    "multi_user_decode":   "multi_user / decode",
}


def _load_sweep_data(sweep_json: Path | str) -> dict:
    sweep_json = Path(sweep_json)
    if not sweep_json.exists():
        return {"rows": [], "config": {}, "panels": []}
    return json.loads(sweep_json.read_text())


def _row_for(rows: list, panel: str) -> dict | None:
    for r in rows:
        if r.get("panel") == panel:
            return r
    return None


def emit_panel_op_log_summary(
    panel: str,
    sweep_json: Path | str = _SWEEP_JSON,
    out_dir: Path | str = _OUTPUT_DIR,
) -> str | None:
    """One bar chart of the 5 op_log counts for ``panel``.

    Returns the written PNG path, or ``None`` when sweep.json is empty
    or the requested panel is absent.
    """
    import matplotlib.pyplot as plt

    data = _load_sweep_data(sweep_json)
    row = _row_for(data.get("rows", []), panel)
    if row is None:
        return None
    summary = row.get("op_log_summary", {})
    n_ranks = row.get("n_ranks")

    labels = [_OP_LOG_DISPLAY[k] for k in _OP_LOG_KEYS]
    values = [summary.get(k, 0) for k in _OP_LOG_KEYS]
    colors = [_OP_LOG_COLORS[k] for k in _OP_LOG_KEYS]

    fig, ax = plt.subplots(figsize=(8, 5))
    bars = ax.bar(labels, values, color=colors)
    for b, v in zip(bars, values):
        ax.text(b.get_x() + b.get_width() / 2, b.get_height(),
                f"{int(v)}", ha="center", va="bottom", fontsize=9)
    ax.set_title(
        f"{_PANEL_DISPLAY.get(panel, panel)}  (n_ranks={n_ranks})",
        fontsize=12, fontweight="bold",
    )
    ax.set_ylabel("count")
    ax.grid(True, axis="y", alpha=0.3)
    fig.tight_layout()

    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    out = out_dir / f"gqa_op_log_{panel}.png"
    fig.savefig(out, dpi=120)
    plt.close(fig)
    return str(out)


def emit_gqa_comparison(
    sweep_json: Path | str = _SWEEP_JSON,
    out_dir: Path | str = _OUTPUT_DIR,
) -> str | None:
    """Grouped-bar chart comparing the 5 op_log counts across all panels."""
    import matplotlib.pyplot as plt
    import numpy as np

    data = _load_sweep_data(sweep_json)
    panels_in = data.get("panels") or list(_PANELS_V1)
    rows = data.get("rows", [])
    panels = [p for p in panels_in if _row_for(rows, p) is not None]
    if not panels:
        return None

    n_groups = len(panels)
    n_series = len(_OP_LOG_KEYS)
    x = np.arange(n_groups)
    width = 0.8 / n_series

    fig, ax = plt.subplots(figsize=(11, 6))
    for i, key in enumerate(_OP_LOG_KEYS):
        offset = (i - (n_series - 1) / 2) * width
        vals = [_row_for(rows, p)["op_log_summary"].get(key, 0)
                for p in panels]
        ax.bar(x + offset, vals, width,
               label=_OP_LOG_DISPLAY[key], color=_OP_LOG_COLORS[key])

    ax.set_xticks(x)
    ax.set_xticklabels(
        [f"{_PANEL_DISPLAY.get(p, p)}\n(n_ranks={_row_for(rows, p)['n_ranks']})"
         for p in panels],
        fontsize=8,
    )
    ax.set_ylabel("count")
    ax.set_title("GQA Llama-70B — op_log summary across panels",
                 fontsize=13, fontweight="bold")
    ax.legend(fontsize=8, loc="upper right")
    ax.grid(True, axis="y", alpha=0.3)
    fig.tight_layout()

    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    out = out_dir / "gqa_comparison.png"
    fig.savefig(out, dpi=120)
    plt.close(fig)
    return str(out)


def emit_all_gqa_plots(
    sweep_json: Path | str = _SWEEP_JSON,
    out_dir: Path | str = _OUTPUT_DIR,
) -> list[str]:
    """Render all 5 in-scope figures and return the written paths.

    Sub-cycle 4c v1 emits 5 of the 6 figures ADR-0057 D3 lists; the
    6th (gqa_scaling.png) needs sub-cycle 4b's Q/cube sweep data.
    """
    paths: list[str] = []
    for panel in _PANELS_V1:
        p = emit_panel_op_log_summary(panel, sweep_json, out_dir)
        if p is not None:
            paths.append(p)
    comp = emit_gqa_comparison(sweep_json, out_dir)
    if comp is not None:
        paths.append(comp)
    return paths


# ── Bench entry ────────────────────────────────────────────────────────


@bench(
    name="milestone-gqa-llama70b",
    description="1H milestone: GQA Llama-70B 4-panel sweep (ADR-0057 v1).",
)
def run(torch) -> None:
    """Drive the four GQA panels at validation scale; write sweep.json and figures.

    Modes (mutually exclusive):
      MILESTONE_FAST=1  Skip the sweep; re-render figures from the
                        committed sweep.json. Seconds, no simulator.
      GQA_VALIDATION=1  Run the four-panel validation sweep + figures.
                        ~1-2h on the full simulator.

    Headline-scale mode is deferred to sub-cycle 4c (figures landed
    here; headline-scale + scaling figure await sub-cycle 4b).
    A sentinel tensor is submitted at the end so run_bench's ADR-0045 D4
    "at least one request" contract is satisfied even when the panels
    are skipped via MILESTONE_FAST=1.
    """
    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    fast = bool(os.environ.get("MILESTONE_FAST"))
    if not fast and not os.environ.get("GQA_VALIDATION"):
        raise RuntimeError(
            "milestone-gqa-llama70b v1 needs GQA_VALIDATION=1 (run the "
            "sweep) or MILESTONE_FAST=1 (reuse committed sweep.json). "
            "Headline mode is deferred to sub-cycle 4b/4c per ADR-0057 D3."
        )

    if not fast:
        topology = os.environ.get("GQA_TOPOLOGY", "topology.yaml")
        rows = [_run_panel(panel, topology) for panel in _PANELS_V1]
        sweep = {
            "version": 1,
            "validation_scale": True,
            "panels": list(_PANELS_V1),
            "config": {
                "S_q_prefill": _S_Q_PREFILL,
                "S_kv_per_rank": _S_KV_PER_RANK,
                "h_q": _H_Q,
                "h_kv": _H_KV,
                "d_head": _D_HEAD,
                "n_ranks_single_user": _N_RANKS_SINGLE_USER,
                "n_ranks_multi_user": _N_RANKS_MULTI_USER,
            },
            "rows": rows,
        }
        _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
        print(f"  milestone-gqa-llama70b: {len(rows)} rows -> {_SWEEP_JSON}")
    elif not _SWEEP_JSON.exists():
        raise RuntimeError(
            f"MILESTONE_FAST=1 requires {_SWEEP_JSON} to exist; "
            "run with GQA_VALIDATION=1 once to seed it."
        )

    paths = emit_all_gqa_plots()
    print(f"  milestone-gqa-llama70b: {len(paths)} figures -> {_OUTPUT_DIR} "
          f"(fast={fast})")

    # Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out).
    torch.zeros(
        (1, 1), dtype="f16",
        dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1),
        name="milestone_gqa_sentinel",
    )
