"""Batch-scaling figure for concurrent decode users on one SIP.

Reads ``docs/sweeps/decode_batch_sweep.csv`` (from
``tests/attention/test_gqa_decode_batch_sweep.py``) and plots aggregate
throughput (requests / us) versus batch size B for the four KV mappings,
one panel per context length. Each mapping's curve runs up to its SIP
capacity 16/C (A1=2 ... B=16 users).
"""
from pathlib import Path

import csv

import matplotlib.pyplot as plt

ROOT = Path(__file__).resolve().parents[3]
CSV = ROOT / "docs" / "sweeps" / "decode_batch_sweep.csv"
OUT = (ROOT / "docs" / "report" / "1H-codesign-paper"
       / "figures" / "gqa_short_context" / "batch_scaling.png")

MODE_LABEL = {"A1": "1-kv-per-cube", "A2": "2-kv-per-cube",
              "A4": "4-kv-per-cube", "B": "8-kv-per-cube"}
MODE_COLOR = {"A1": "#1f77b4", "A2": "#2ca02c",
              "A4": "#ffd000", "B": "#d62728"}
MODES = ["A1", "A2", "A4", "B"]


def main():
    rows = list(csv.DictReader(CSV.open()))
    contexts = sorted({int(r["S_kv"]) for r in rows})
    fig, axes = plt.subplots(1, len(contexts), figsize=(6.5 * len(contexts), 4.5),
                             squeeze=False)
    xmax = max(int(r["B"]) for r in rows)  # largest SIP capacity (= 16)
    for col, S in enumerate(contexts):
        ax = axes[0][col]
        for m in MODES:
            pts = sorted(
                ((int(r["B"]), float(r["throughput_users_per_us"]))
                 for r in rows if r["mode"] == m and int(r["S_kv"]) == S),
                key=lambda t: t[0],
            )
            if not pts:
                continue
            xs, ys = zip(*pts)
            # Measured: concurrent users up to the SIP capacity 16/C.
            ax.plot(xs, ys, marker="o", color=MODE_COLOR[m],
                    label=MODE_LABEL[m])
            # Beyond capacity the SIP is full: extra users run in waves at
            # the saturated rate, so throughput plateaus. Draw that ceiling
            # as a dashed extension (projected, not concurrently measured).
            if xs[-1] < xmax:
                ax.plot([xs[-1], xmax], [ys[-1], ys[-1]],
                        linestyle="--", color=MODE_COLOR[m], alpha=0.55)
            ax.annotate(f"{ys[-1]:.2f}", (xs[-1], ys[-1]),
                        textcoords="offset points", xytext=(4, 4), fontsize=8)
        ax.set_title(f"S_kv = {S // 1024}K")
        ax.set_xlabel("batch size B (concurrent users)")
        ax.set_ylabel("throughput (requests / µs)")
        ax.grid(True, alpha=0.3)
        ax.legend(fontsize=9)
    fig.suptitle("Batch scaling: aggregate throughput vs concurrent users "
                 "(one SIP)", fontsize=12, fontweight="bold")
    fig.tight_layout()
    OUT.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(OUT, dpi=140, bbox_inches="tight")
    print(f"  ✓ {OUT}")


if __name__ == "__main__":
    main()