"""Comparative figure for the Case-6 composite-command decode study.

Reads sweep_decode_composite.json (emitted by milestone-1h-gqa, sweep
``composite``) and writes one two-panel PNG into the bench-output dir:

  gqa_decode_long_ctx_composite.png
    Left  — end-to-end decode latency (µs) vs context length, per command
            form (primitive / composite / composite_extended).
    Right — PE_CPU command count vs context length: the hand-tiled
            primitive kernel issues O(n_tiles) commands (rises with
            context), while the coarse composite forms issue O(1) and
            *saturate* — PE_SCHEDULER absorbs the per-tile fan-out.

The x-axis is the global context length S_kv; each PE owns
S_local = S_kv/(C·P=64) tokens, so at the 1M production point every PE
runs Q·Kᵀ of (G·T_q, d_head)·(d_head, 16384) and P·V of
(G·T_q, 16384)·(16384, d_head).

Run (after the bench):
    GQA_1H_RUN=1 GQA_1H_SWEEPS=composite python -m kernbench.cli.main run \\
        --bench milestone-1h-gqa --topology topology.yaml
    python scripts/paper/paper_plot_gqa_decode_long_ctx_composite.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 / "src" / "kernbench" / "benches"
    / "1H_milestone_output" / "gqa" / "long_ctx"
)
_SWEEP_JSON = _FIG_DIR / "sweep_decode_composite.json"
_PAPER_FIG_DIR = (
    _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
)

_N_RANKS = 64  # C·P for the Case-6 64-way split.

# variant key → (display label, colour, marker)
_VARIANT_STYLE = {
    "primitive":          ("primitive (tl.dot, hand-tiled)", "#c0504d", "o"),
    "composite":          ("composite GEMM",                 "#3b6ea5", "s"),
    "composite_extended": ("composite + softmax_merge",      "#4f8a4f", "^"),
}
_ORDER = ("primitive", "composite", "composite_extended")


def _load() -> dict:
    return json.loads(_SWEEP_JSON.read_text())


def _series(rows: list[dict], variant: str, key: str):
    """Sorted (S_kv, value) series for a variant, skipping null values
    (latency is only measured over the tractable S_kv subset)."""
    pts = sorted(
        ((r["S_kv"], r[key]) for r in rows
         if r["variant"] == variant and r.get(key) is not None),
        key=lambda t: t[0],
    )
    return [p[0] for p in pts], [p[1] for p in pts]


def _xticklabels(s_kvs: list[int]) -> list[str]:
    out = []
    for s in s_kvs:
        if s >= 1 << 20:
            out.append(f"{s // (1 << 20)}M")
        else:
            out.append(f"{s // 1024}K")
    return out


def main() -> None:
    sweep = _load()
    rows = sweep["rows"]
    s_kv_op = sweep["s_kv_opcount"]
    s_kv_lat = sweep["s_kv_latency"]

    fig, (ax_lat, ax_cmd) = plt.subplots(1, 2, figsize=(13.0, 4.8))

    for v in _ORDER:
        label, color, marker = _VARIANT_STYLE[v]
        xs, lat = _series(rows, v, "latency_ns")
        ax_lat.plot(xs, [y / 1e3 for y in lat], marker=marker,
                    color=color, label=label, lw=2)
        xs, cmds = _series(rows, v, "pe_cpu_cmd_count")
        ax_cmd.plot(xs, cmds, marker=marker, color=color, label=label, lw=2)

    ax_lat.set_xticks(s_kv_lat)
    ax_lat.set_xticklabels(_xticklabels(s_kv_lat))
    ax_cmd.set_xticks(s_kv_op)
    ax_cmd.set_xticklabels(_xticklabels(s_kv_op), fontsize=8)
    for ax in (ax_lat, ax_cmd):
        ax.set_xscale("log", base=2)
        ax.set_xlabel(
            r"context length $S_{kv}$ "
            r"($S_{\mathrm{local}}=S_{kv}/64$ per PE)"
        )
        ax.grid(True, ls=":", alpha=0.5)
        ax.legend(fontsize=9)

    ax_lat.set_ylabel("end-to-end decode latency (µs)")
    ax_lat.set_title(
        "Case-6 decode latency per command form\n"
        "(memory-bound: command form does not move the critical path)"
    )
    ax_cmd.set_ylabel("PE_CPU commands issued")
    ax_cmd.set_title(
        "PE_CPU command count per command form\n"
        "(primitive O(n$_\\mathrm{tiles}$) rises; composite O(1) saturates)"
    )
    ax_cmd.axvline(1 << 20, color="#888", ls="--", lw=1, alpha=0.7)
    ax_cmd.annotate("1M production\ncontext", xy=(1 << 20, 0),
                    xytext=(1 << 18, 0.6), fontsize=8,
                    textcoords=("data", "axes fraction"), ha="right",
                    color="#555")

    fig.suptitle(
        "Case-6 (Cube-SP × PE-SP) long-context decode — use of composite "
        "commands\nLLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs), "
        "$T_q{=}1$",
        fontsize=11,
    )
    fig.tight_layout(rect=(0, 0, 1, 0.94))

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

    # Mirror into the paper figures dir (derived artifact).
    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()