"""6-case KV-sharding tensor diagram (the slide-13 PNG export).

Flat 2-D rectangles, one per sharding case, with:
    Y axis = S_kv (vertical)        — Cube-SP / PE-SP slice it
    X axis = d_head (horizontal)    — Cube-TP-d_head / PE-TP-d_head slice it

Drops the batch axis entirely (decode: B = 1, T_q = 1). Same case set
and visual encoding as slide 13 of GQA_full_deck.pptx; matplotlib
renders it cleanly so the PNG sits next to the other GQA summary
artifacts in 1H_milestone_output/gqa/long_ctx/.
"""
from __future__ import annotations

from pathlib import Path

import matplotlib.patches as mpatches
import matplotlib.pyplot as plt

_C = 8
_P = 8

_GROUP_FILLS = [
    "#A5D8FF", "#B2F2BB", "#FFD8A8", "#FFC9C9",
    "#D0BFFF", "#99E9F2", "#FCC2D7", "#FFEC99",
]

_ACC = {
    "red":    "#E03131",
    "orange": "#FD7E14",
    "blue":   "#1C7ED6",
    "green":  "#37B24D",
}

# (label, accent, kv, comm, overflow, encoding-flags, axis-spec)
#   y_split   = 8 horizontal Y bands  (Cube-SP on S_kv)
#   x_split   = 8 vertical X bands    (Cube-TP-d_head)
#   pe_y      = 7 fine horizontal dividers within each Y band
#   pe_x      = 7 fine vertical dividers within each X band
#   axes      = small annotation under the chip naming the axes
#               that the cube/PE actually shard, so the reader can
#               parse Case 5 (where cube colour fills run X instead
#               of Y, breaking the visual symmetry of the rest).
_CASES = [
    dict(label="Case 1\nCube-Repl / PE-repl", accent=_ACC["red"],
         kv="40 GB", comm="1.2 MB", overflow=True,
         y_split=False, x_split=False, pe_y=False, pe_x=False,
         axes="Cube: replicated     PE: replicated"),
    dict(label="Case 2\nCube-SP / PE-repl",   accent=_ACC["orange"],
         kv="5 GB",  comm="3.8 MB", overflow=True,
         y_split=True,  x_split=False, pe_y=False, pe_x=False,
         axes="Cube → Y (S_kv)     PE: replicated"),
    dict(label="Case 3\nCube-Repl / PE-SP",   accent=_ACC["orange"],
         kv="5 GB",  comm="3.8 MB", overflow=True,
         y_split=False, x_split=False, pe_y=True,  pe_x=False,
         axes="Cube: replicated     PE → Y (S_kv)"),
    dict(label="Case 4\nCube-SP / PE-TP-d_head", accent=_ACC["blue"],
         kv="640 MB", comm="166 MB", overflow=False,
         y_split=True,  x_split=False, pe_y=False, pe_x=True,
         axes="Cube → Y (S_kv)     PE → X (d_head)"),
    dict(label="Case 5\nCube-TP-d_head / PE-SP", accent=_ACC["blue"],
         kv="640 MB", comm="166 MB", overflow=False,
         y_split=False, x_split=True,  pe_y=True,  pe_x=False,
         axes="Cube → X (d_head)     PE → Y (S_kv)"),
    dict(label="Case 6 ★\nCube-SP / PE-SP",   accent=_ACC["green"],
         kv="640 MB", comm="6.2 MB", overflow=False,
         y_split=True,  x_split=False, pe_y=True,  pe_x=False,
         axes="Cube → Y (S_kv)     PE → Y (S_kv)"),
]

_OUT_DIR = (
    Path(__file__).resolve().parents[2]
    / "src" / "kernbench" / "benches"
    / "1H_milestone_output" / "gqa" / "long_ctx"
)


def _draw_panel(ax, cfg):
    """Draw one case's 2-D KV-tensor rectangle into a panel ax."""
    ax.set_xlim(0, 1)
    ax.set_ylim(1, 0)        # Y points down (S_kv ↓)
    ax.set_aspect("auto")
    ax.set_xticks([])
    ax.set_yticks([])

    cube_repl = not cfg["y_split"] and not cfg["x_split"]
    pe_repl   = not cfg["pe_y"]   and not cfg["pe_x"]

    # Cube-level colour fill.
    if cfg["y_split"] and not cfg["x_split"]:
        # 8 horizontal Y bands.
        for c in range(_C):
            ax.add_patch(mpatches.Rectangle(
                (0, c / _C), 1, 1 / _C,
                facecolor=_GROUP_FILLS[c], edgecolor="black", linewidth=0.6))
            ax.text(0.04, c / _C + 0.5 / _C, f"C{c}",
                    ha="left", va="center", fontsize=8,
                    fontweight="bold", color="#333")
    elif cfg["x_split"] and not cfg["y_split"]:
        # 8 vertical X bands.
        for c in range(_C):
            ax.add_patch(mpatches.Rectangle(
                (c / _C, 0), 1 / _C, 1,
                facecolor=_GROUP_FILLS[c], edgecolor="black", linewidth=0.6))
            ax.text(c / _C + 0.5 / _C, 0.04, f"C{c}",
                    ha="center", va="top", fontsize=8,
                    fontweight="bold", color="#333")
    else:
        ax.add_patch(mpatches.Rectangle(
            (0, 0), 1, 1,
            facecolor="#F5F5F5", edgecolor="black", linewidth=0.8))
        ax.text(0.5, 0.5, "× 8 cubes\nfull KV",
                ha="center", va="center",
                fontsize=10, fontweight="bold",
                fontstyle="italic", color="#666")

    # PE-level fine dividers — distinguished from cube boundaries by
    # using a dashed style + slightly stronger contrast. This is what
    # makes Case 5's PE-SP (horizontal lines across vertical cube
    # bands) read as "different axis from the cubes" at a glance.
    if cfg["pe_y"]:
        outer = _C if cfg["y_split"] else 1
        band = 1 / outer
        for o in range(outer):
            for p in range(1, _P):
                y = o * band + band * p / _P
                ax.axhline(y, color="#222", linewidth=0.8,
                           linestyle=(0, (3, 2)), alpha=0.75)
    if cfg["pe_x"]:
        outer = _C if cfg["x_split"] else 1
        band = 1 / outer
        for o in range(outer):
            for p in range(1, _P):
                x = o * band + band * p / _P
                ax.axvline(x, color="#222", linewidth=0.8,
                           linestyle=(0, (3, 2)), alpha=0.75)

    # Heavy outline on top.
    ax.add_patch(mpatches.Rectangle(
        (0, 0), 1, 1, facecolor="none",
        edgecolor="black", linewidth=1.2))

    # Replication badges — small text-only badges in the corners of
    # the rectangle, no ghost-card stacking (which mis-reads as a
    # larger enclosing tensor).
    badges: list[str] = []
    if cube_repl:
        badges.append("× 8 cube copies")
    if pe_repl and (cfg["y_split"] or cfg["x_split"]):
        # Cube is sharded but PEs in each cube replicate that shard.
        badges.append("× 8 PEs / cube replicate")
    elif pe_repl and cube_repl:
        # Both replicated — PE replication adds to the cube one.
        badges.append("× 8 PEs / cube replicate")
    if badges:
        ax.text(0.98, 0.02, "\n".join(badges),
                ha="right", va="top", fontsize=7,
                fontweight="bold", color="#444",
                fontstyle="italic",
                bbox=dict(facecolor="white", edgecolor="#888",
                          boxstyle="round,pad=0.20", linewidth=0.5))


def _make_table_png() -> Path:
    """Slide-14 companion table: per-PE memory + comm for all 6 cases."""
    headers = ["Case", "Sharding", "KV / PE", "Fit",
               "Comm/tok\n(analytical)", "Comm/tok\n(measured)", "Notes"]
    rows = [
        ("Case 1",   "Cube-Repl  ·  PE-repl",      "40 GB",  "✗",
         "1.2 MB",  "1.25 MB",
         "no sharding — full KV on every PE"),
        ("Case 2",   "Cube-SP    ·  PE-repl",      "5 GB",   "✗",
         "3.8 MB",  "1.27 MB",
         "cube-axis sharded only"),
        ("Case 3",   "Cube-Repl  ·  PE-SP",        "5 GB",   "✗",
         "3.8 MB",  "1.39 MB",
         "PE-axis sharded only"),
        ("Case 4",   "Cube-SP    ·  PE-TP-d_head", "640 MB", "✓",
         "166 MB",  "162 MB",
         "d_head split intra-cube — partial-score AR ∝ S_kv"),
        ("Case 5",   "Cube-TP-d_head  ·  PE-SP",   "640 MB", "✓",
         "166 MB",  "162 MB",
         "d_head split inter-cube — partial-score AR on UCIe"),
        ("Case 6 ★", "Cube-SP    ·  PE-SP",        "640 MB", "✓",
         "6.2 MB",  "1.41 MB",
         "S_kv split on both axes — (m,ℓ,O) AR only, S_kv-indep."),
    ]
    accents = [_ACC["red"], _ACC["orange"], _ACC["orange"],
               _ACC["blue"], _ACC["blue"], _ACC["green"]]

    fig, ax = plt.subplots(figsize=(20.0, 4.6))
    ax.set_axis_off()

    cell_data = [headers] + [list(r) for r in rows]
    tbl = ax.table(cellText=cell_data,
                   colWidths=[0.07, 0.18, 0.08, 0.05, 0.12, 0.12, 0.38],
                   cellLoc="center", loc="center")
    tbl.auto_set_font_size(False)
    tbl.set_fontsize(11)
    tbl.scale(1.0, 2.2)

    n_cols = len(headers)
    n_rows = len(rows) + 1   # +1 header
    # Header styling.
    for ci in range(n_cols):
        cell = tbl[(0, ci)]
        cell.set_facecolor("#1F4E79")
        cell.set_text_props(color="white", weight="bold")
        cell.set_edgecolor("#1F4E79")
    # Body styling.
    for ri, row in enumerate(rows, start=1):
        is_pareto = row[0].endswith("★")
        row_fill = "#E8F5E9" if is_pareto else (
            "white" if ri % 2 == 1 else "#F5F5F7")
        # Case-name cell uses accent.
        case_cell = tbl[(ri, 0)]
        case_cell.set_facecolor(accents[ri - 1])
        case_cell.set_text_props(color="white", weight="bold")
        # Remaining cells.
        for ci in range(1, n_cols):
            cell = tbl[(ri, ci)]
            cell.set_facecolor(row_fill)
            txt_kwargs = {"weight": "bold" if is_pareto else "normal",
                          "color": "#333"}
            if ci == 2:                              # KV / PE
                txt_kwargs["color"] = (
                    "#C62828" if row[3] == "✗" else "#2E7D32")
                txt_kwargs["weight"] = "bold"
            if ci == 3:                              # Fit
                txt_kwargs["color"] = (
                    "#C62828" if row[3] == "✗" else "#2E7D32")
                txt_kwargs["weight"] = "bold"
            cell.set_text_props(**txt_kwargs)
        # Last-column (Notes) cells left-aligned for readability.
        tbl[(ri, n_cols - 1)].get_text().set_ha("left")

    # Force left-align on the Notes header too.
    tbl[(0, n_cols - 1)].get_text().set_ha("left")

    fig.suptitle(
        "GQA decode KV-sharding — per-PE memory & communication  "
        "(LLaMA 70B GQA single KV-head group · S_kv = 1 M, FP16, 80 layers)",
        fontsize=12, y=0.97,
    )
    out = _OUT_DIR / "gqa_kv_sharding_6cases_table.png"
    fig.savefig(out, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {out}")
    return out


def main() -> Path:
    _OUT_DIR.mkdir(parents=True, exist_ok=True)
    n = len(_CASES)
    fig = plt.figure(figsize=(20.0, 7.0))
    # Three rows per column: case chip · axis-spec annotation · rectangle.
    gs = fig.add_gridspec(3, n,
                          height_ratios=[0.55, 0.32, 8.5],
                          hspace=0.05, wspace=0.20,
                          left=0.04, right=0.99,
                          top=0.93, bottom=0.06)

    for i, cfg in enumerate(_CASES):
        # Top: case chip header.
        ax_chip = fig.add_subplot(gs[0, i])
        ax_chip.set_xticks([])
        ax_chip.set_yticks([])
        for spine in ax_chip.spines.values():
            spine.set_visible(False)
        ax_chip.add_patch(mpatches.Rectangle(
            (0, 0), 1, 1, transform=ax_chip.transAxes,
            facecolor=cfg["accent"], edgecolor=cfg["accent"]))
        ax_chip.text(0.5, 0.5, cfg["label"],
                     ha="center", va="center",
                     fontsize=10, fontweight="bold",
                     color="white")

        # Middle: axis-spec annotation — names which axis the cube
        # shards on and which axis the PE shards on (essential for
        # parsing Case 5 where the cube colour fills run X instead
        # of Y, breaking the visual symmetry of the rest).
        ax_axes = fig.add_subplot(gs[1, i])
        ax_axes.set_xticks([])
        ax_axes.set_yticks([])
        for spine in ax_axes.spines.values():
            spine.set_visible(False)
        ax_axes.add_patch(mpatches.Rectangle(
            (0, 0), 1, 1, transform=ax_axes.transAxes,
            facecolor="#F5F5F7", edgecolor="#CCCCCC",
            linewidth=0.6))
        ax_axes.text(0.5, 0.5, cfg["axes"],
                     ha="center", va="center",
                     fontsize=8.5, fontweight="bold",
                     color="#1F4E79")

        # Bottom: the tensor rectangle.
        ax = fig.add_subplot(gs[2, i])
        _draw_panel(ax, cfg)
        ax.set_xlabel("X : d_head = 128  →",
                      fontsize=9, fontweight="bold",
                      fontstyle="italic", color="#1F4E79")
        ax.set_ylabel("Y : S_kv = 1 M  ↓",
                      fontsize=9, fontweight="bold",
                      fontstyle="italic", color="#1F4E79")

    fig.suptitle(
        "GQA decode KV-tensor sharding — 6 cases  ·  "
        "LLaMA 70B GQA single KV-head group  ·  "
        "C = 8 cubes × P = 8 PEs  ·  S_kv = 1 M, FP16, 80 layers",
        fontsize=12, y=0.99,
    )

    out = _OUT_DIR / "gqa_kv_sharding_6cases_diagram.png"
    fig.savefig(out, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {out}")

    # Companion table PNG (slide-14 export).
    _make_table_png()
    return out


if __name__ == "__main__":
    main()