"""Generate 6 comparison figures for the short-context GQA attention sweep.

Sweep matrix: 3 variants × 2 phases × 4 modes × 4 contexts (8K/16K/32K/64K).

Figures
  Wall-clock latency:
    1. wall_variant1_mode_compare.png  — (1) without composite
    2. wall_variant2_mode_compare.png  — (2) with composite (GEMM-only)
    3. wall_variant3_mode_compare.png  — (3) with composite + softmax_merge
    4. wall_a1_variant_compare.png     — 1-kv-per-cube: variant comparison

  Mode trade-off + composite ablation:
    5. per_cube_tradeoff_mode_compare.png      — HBM BW + IPCQ (variant-invariant)
    6. gemm_util_a1_variant_ablation.png       — 1-kv-per-cube GEMM util

Output: docs/report/1H-codesign-paper/figures/gqa_short_context/
"""
from __future__ import annotations

import csv
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

ROOT = Path(__file__).resolve().parents[3]
SWEEPS = ROOT / "docs" / "sweeps"
OUT = (ROOT / "docs" / "report" / "1H-codesign-paper"
       / "figures" / "gqa_short_context")
OUT.mkdir(parents=True, exist_ok=True)

CONTEXTS = [8192, 16384, 32768, 65536]
CONTEXT_LABELS = ["8K", "16K", "32K", "64K"]
MODES = ["A1", "A2", "A4", "B"]
MODE_LABEL = {
    "A1": "1-kv-per-cube",
    "A2": "2-kv-per-cube",
    "A4": "4-kv-per-cube",
    "B":  "8-kv-per-cube",
}
MODE_COLOR = {
    MODE_LABEL["A1"]: "#1f77b4",
    MODE_LABEL["A2"]: "#2ca02c",
    MODE_LABEL["A4"]: "#ffd000",
    MODE_LABEL["B"]:  "#d62728",
}
VARIANTS = [
    ("baseline",        "(1) without composite"),
    ("composite",       "(2) with composite (GEMM-only)"),
    ("composite_fused", "(3) with composite + softmax_merge"),
]
VARIANT_COLOR = {
    "baseline":        "#1f77b4",
    "composite":       "#2ca02c",
    "composite_fused": "#d62728",
}

CSV_MAP = {
    ("decode",  "baseline"):         SWEEPS / "short_context_decode_sweep.csv",
    ("decode",  "composite"):        SWEEPS / "short_context_decode_composite_sweep.csv",
    ("decode",  "composite_fused"):  SWEEPS / "short_context_decode_composite_fused_sweep.csv",
    ("prefill", "baseline"):         SWEEPS / "short_context_prefill_sweep.csv",
    ("prefill", "composite"):        SWEEPS / "short_context_prefill_composite_sweep.csv",
    ("prefill", "composite_fused"):  SWEEPS / "short_context_prefill_composite_fused_sweep.csv",
}


def load(phase, variant):
    """Return {(mode, S_kv): row dict}."""
    rows = {}
    with CSV_MAP[(phase, variant)].open() as f:
        for r in csv.DictReader(f):
            rows[(r["mode"], int(r["S_kv"]))] = r
    return rows


def grouped_bars(ax, *, x_labels, groups, group_color, ylabel, title,
                 y_log=False, value_fmt="{:.3g}"):
    """Grouped bars: x = x_labels, groups = list of (label, values)."""
    n = len(groups)
    nx = len(x_labels)
    width = 0.8 / n
    x = np.arange(nx)
    for i, (label, values) in enumerate(groups):
        offset = (i - (n - 1) / 2) * width
        bars = ax.bar(x + offset, values, width, label=label,
                      color=group_color[label], edgecolor="black",
                      linewidth=0.4)
        for bar, v in zip(bars, values):
            if v <= 0 and y_log:
                continue
            ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height(),
                    value_fmt.format(v), ha="center", va="bottom",
                    fontsize=7)
    ax.set_xticks(x)
    ax.set_xticklabels(x_labels)
    ax.set_xlabel("S_kv")
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    if y_log:
        ax.set_yscale("log")
    ax.grid(True, axis="y", alpha=0.3, which="both")
    ax.legend(fontsize=8, loc="upper left")


# ── 1-3. Wall mode comparison per variant ──────────────────────────

def fig_wall_mode_compare(variant_key, suptitle, fname):
    fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
    for ax, phase in zip(axes, ("prefill", "decode")):
        rows = load(phase, variant_key)
        groups = [
            (MODE_LABEL[m], [float(rows[(m, s)]["wall_us"]) for s in CONTEXTS])
            for m in MODES
        ]
        grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
                     group_color=MODE_COLOR,
                     ylabel="Wall-clock latency (μs)",
                     title=phase.capitalize(),
                     y_log=True, value_fmt="{:.0f}")
    fig.suptitle(suptitle, fontsize=12, fontweight="bold")
    fig.tight_layout()
    out = OUT / fname
    fig.savefig(out, dpi=140, bbox_inches="tight")
    plt.close(fig)
    print(f"  ✓ {out}")


# ── 4. Wall 1-kv-per-cube variant comparison ───────────────────────

def fig_wall_a1_variant_compare(fname):
    fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
    for ax, phase in zip(axes, ("prefill", "decode")):
        groups = []
        for variant_key, _ in VARIANTS:
            rows = load(phase, variant_key)
            values = [float(rows[("A1", s)]["wall_us"]) for s in CONTEXTS]
            groups.append((variant_key, values))
        grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
                     group_color=VARIANT_COLOR,
                     ylabel="Wall-clock latency (μs)",
                     title=phase.capitalize(),
                     y_log=True, value_fmt="{:.0f}")
        ax.legend([v[1] for v in VARIANTS], fontsize=8, loc="upper left")
    fig.suptitle("1-kv-per-cube: variant comparison",
                 fontsize=12, fontweight="bold")
    fig.tight_layout()
    out = OUT / fname
    fig.savefig(out, dpi=140, bbox_inches="tight")
    plt.close(fig)
    print(f"  ✓ {out}")


# ── 5. Per-cube trade-off (HBM BW + IPCQ, mode compare) ────────────

def fig_per_cube_tradeoff(fname):
    fig, axes = plt.subplots(2, 2, figsize=(13, 9))
    metrics = [
        ("kv_cache_per_cube_mb", "Per-cube KV footprint (MB)", True, "{:.0f}"),
        ("ipcq_kb",              "IPCQ traffic (KB)",          True, "{:.1f}"),
    ]
    for row, (key, ylabel, log, fmt) in enumerate(metrics):
        for col, phase in enumerate(("prefill", "decode")):
            ax = axes[row][col]
            rows = load(phase, "baseline")
            groups = [
                (MODE_LABEL[m], [float(rows[(m, s)][key]) for s in CONTEXTS])
                for m in MODES
            ]
            grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
                         group_color=MODE_COLOR, ylabel=ylabel,
                         title=phase.capitalize(),
                         y_log=log, value_fmt=fmt)
    fig.suptitle("Per-cube trade-off: KV footprint + IPCQ mode comparison",
                 fontsize=12, fontweight="bold")
    fig.tight_layout()
    out = OUT / fname
    fig.savefig(out, dpi=140, bbox_inches="tight")
    plt.close(fig)
    print(f"  ✓ {out}")


# ── 6. GEMM util ablation (1-kv-per-cube, variant compare) ─────────

def fig_gemm_util_a1_ablation(fname):
    fig, axes = plt.subplots(1, 2, figsize=(13, 5.4))
    for ax, phase in zip(axes, ("prefill", "decode")):
        groups = []
        for variant_key, _ in VARIANTS:
            rows = load(phase, variant_key)
            values = [float(rows[("A1", s)]["gemm_util"]) for s in CONTEXTS]
            groups.append((variant_key, values))
        grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
                     group_color=VARIANT_COLOR,
                     ylabel="GEMM engine utilization (per PE)",
                     title=phase.capitalize(),
                     y_log=False, value_fmt="{:.4f}")
        ax.legend([v[1] for v in VARIANTS], fontsize=8, loc="upper left")
    fig.suptitle("1-kv-per-cube: GEMM utilization variant comparison",
                 fontsize=12, fontweight="bold")
    fig.text(
        0.5, -0.04,
        "⚠ Higher gemm_util on composite/fused is largely supertile "
        "padding overhead (M=G=8 padded to TILE_M=32), not pure fusion "
        "gain. See ADR-0070 limitation #4.",
        ha="center", fontsize=9, style="italic", color="#555")
    fig.tight_layout()
    out = OUT / fname
    fig.savefig(out, dpi=140, bbox_inches="tight")
    plt.close(fig)
    print(f"  ✓ {out}")


if __name__ == "__main__":
    print(f"Output: {OUT}")
    fig_wall_mode_compare(
        "baseline",
        "(1) Without composite: latency mode comparison",
        "wall_variant1_mode_compare.png")
    fig_wall_mode_compare(
        "composite",
        "(2) With composite (GEMM-only): latency mode comparison",
        "wall_variant2_mode_compare.png")
    fig_wall_mode_compare(
        "composite_fused",
        "(3) With composite + softmax_merge: latency mode comparison",
        "wall_variant3_mode_compare.png")
    fig_wall_a1_variant_compare("wall_a1_variant_compare.png")
    fig_per_cube_tradeoff("per_cube_tradeoff_mode_compare.png")
    fig_gemm_util_a1_ablation("gemm_util_a1_variant_ablation.png")
    print("\nDone.")
