"""Generate 4 wall-clock comparison figures for the short-context GQA
attention sweep (prefill + decode × A1/A2/A4/B × 3 variants × 4 contexts).

Figures 1-3: mode comparison within one variant
  • 2 panels (prefill, decode), grouped bars per S_kv, hue=mode (A1/A2/A4/B)
Figure 4: A1 variant comparison (best mode)
  • 2 panels, grouped bars per S_kv, hue=variant

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("/Users/lalicorne/Desktop/ahbm/kernbench2")
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): wall_us}."""
    rows = {}
    with CSV_MAP[(phase, variant)].open() as f:
        for r in csv.DictReader(f):
            rows[(r["mode"], int(r["S_kv"]))] = float(r["wall_us"])
    return rows


def _bar_chart(ax, *, x_labels, groups, group_color, ylabel, title,
               y_log=True):
    """Grouped bars: x = x_labels, groups = list of (label, values).

    values is a list aligned with x_labels.
    """
    n_groups = len(groups)
    n_x = len(x_labels)
    width = 0.8 / n_groups
    x = np.arange(n_x)
    for i, (label, values) in enumerate(groups):
        offset = (i - (n_groups - 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):
            ax.text(bar.get_x() + bar.get_width() / 2,
                    bar.get_height(),
                    f"{v:.0f}", ha="center", va="bottom", fontsize=7,
                    rotation=0)
    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")


def _fig_mode_compare(variant_key, variant_label, fname):
    """Figure: 2 panels (prefill/decode), x=S_kv, grouped bars per mode."""
    fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
    for ax, phase in zip(axes, ("prefill", "decode")):
        rows = load(phase, variant_key)
        groups = []
        for mode in MODES:
            values = [rows[(mode, s)] for s in CONTEXTS]
            groups.append((MODE_LABEL[mode], values))
        _bar_chart(ax, x_labels=CONTEXT_LABELS, groups=groups,
                   group_color=MODE_COLOR,
                   ylabel="Wall-clock latency (μs)",
                   title=f"{phase.capitalize()}")
    fig.suptitle(variant_label,
                 fontsize=12, fontweight="bold")
    fig.tight_layout()
    out = OUT / fname
    fig.savefig(out, dpi=140, bbox_inches="tight")
    plt.close(fig)
    print(f"  ✓ {out}")


def _fig_a1_variant_compare(fname):
    """Figure: A1 (best mode), 2 panels, x=S_kv, grouped bars per variant."""
    fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
    for ax, phase in zip(axes, ("prefill", "decode")):
        groups = []
        for variant_key, variant_label in VARIANTS:
            rows = load(phase, variant_key)
            values = [rows[("A1", s)] for s in CONTEXTS]
            groups.append((variant_key, values))
        _bar_chart(
            ax, x_labels=CONTEXT_LABELS, groups=groups,
            group_color=VARIANT_COLOR,
            ylabel="Wall-clock latency (μs)",
            title=f"{phase.capitalize()}",
        )
        # Custom legend labels
        handles, _ = ax.get_legend_handles_labels()
        readable = [v[1] for v in VARIANTS]
        ax.legend(handles, readable, 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}")


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