"""Composite vs async-loading GEMM TFLOPS comparison.

For each shape in the milestone sweep, run both benches:
  - matmul-composite (load_ref variant — A pre-staged, B streamed by
    scheduler inside one composite command)
  - matmul-async (A and B both async-loaded via tl.load, then a single
    tl.dot — no per-tile overlap of streaming-B with GEMM)

Compute per-PE TFLOPS = 2*M*K*N / pe_window_ns for each kernel and emit a
side-by-side bar chart PNG to
  src/kernbench/benches/1H_milestone_output/gemm/gemm_composite_vs_async_tflops.png

Run from repo root:
  python scripts/paper/paper_plot_gemm_async_vs_composite.py
"""
from __future__ import annotations

import json
import os
import sys
from pathlib import Path

REPO = Path(__file__).resolve().parents[2]
OUT_DIR = REPO / "src" / "kernbench" / "benches" / "1H_milestone_output" / "gemm"
OUT_PNG = OUT_DIR / "gemm_composite_vs_async_tflops.png"
OUT_JSON = OUT_DIR / "gemm_composite_vs_async.json"

TOPO_PATH = REPO / "topology.yaml"

# Same 7 shapes as the milestone composite sweep (skip square=512: B alone
# is 512 KiB; async needs all of A+B+out in TCM scratch and the 512×512
# square pushes total scratch use past the 1 MiB cap).
SHAPES: list[tuple[int, int, int]] = [
    (32, 32, 32),     # 1 tile, K<TILE_K (under-tile in K)
    (32, 64, 32),     # 1 tile
    (32, 128, 32),    # 2 tiles
    (32, 128, 128),   # 8 tiles
    (32, 3072, 32),   # 48 tiles (deep-K)
    (8, 128, 128),    # under-tile in M
    (128, 8, 128),    # under-tile in K
]

ENGINES = ("pe_dma", "pe_fetch_store", "pe_gemm", "pe_math")
STAGES = ("DMA_READ", "DMA_WRITE", "FETCH", "STORE", "GEMM", "MATH")


def _pe_records(op_log):
    return [r for r in op_log
            if any(r.component_id.endswith("." + e) for e in ENGINES)]


def _pe_window_ns(op_log) -> float:
    pe = _pe_records(op_log)
    if not pe:
        return 0.0
    return max(r.t_end for r in pe) - min(r.t_start for r in pe)


def _composite_window_ns(op_log) -> float:
    """For the composite kernel: window of records carrying a stage_type
    set by the composite plan (DMA_READ, FETCH, GEMM, STORE, DMA_WRITE).
    Excludes the initial up-front tl.load(A) because that record is an
    atomic DmaReadCmd with no stage_type. Matches the existing
    milestone_1h_gemm.py / gemm_per_pe_tflops.png methodology.
    """
    stage_records = [r for r in op_log
                     if r.params.get("stage_type") in STAGES]
    if not stage_records:
        return 0.0
    return max(r.t_end for r in stage_records) \
           - min(r.t_start for r in stage_records)


def _async_engine_window_ns(op_log) -> float:
    """For the async kernel: engine pipeline window that excludes the
    initial tl.load(A), paralleling composite_window's exclusion of the
    up-front A pre-stage. The first DMA_READ record on pe_dma is the
    tl.load(A); the window starts at the SECOND pe_dma record's t_start
    (= tl.load(B)) and ends at the last engine record's t_end.
    """
    pe = _pe_records(op_log)
    if not pe:
        return 0.0
    dma = sorted(
        (r for r in op_log if r.component_id.endswith(".pe_dma")),
        key=lambda r: r.t_start,
    )
    if len(dma) < 2:
        return _pe_window_ns(op_log)
    window_start = dma[1].t_start
    window_end = max(r.t_end for r in pe)
    return window_end - window_start


def _run_one(bench_name: str, variant: str | None, M: int, K: int, N: int) -> dict:
    os.environ["MATMUL_M"] = str(M)
    os.environ["MATMUL_K"] = str(K)
    os.environ["MATMUL_N"] = str(N)
    if variant is not None:
        os.environ["MATMUL_VARIANT"] = variant
    elif "MATMUL_VARIANT" in os.environ:
        del os.environ["MATMUL_VARIANT"]

    from kernbench.benches.registry import resolve as resolve_bench
    from kernbench.runtime_api.bench_runner import run_bench
    from kernbench.runtime_api.types import resolve_device
    from kernbench.sim_engine.engine import GraphEngine
    from kernbench.topology.builder import resolve_topology

    # Chunked async loads A in TILE_K-sized blocks at sequential offsets;
    # those offsets do not match the row-major (M, K) layout of A, so the
    # DataExecutor would fail on the resulting mismatched read. The
    # simulator's *timing* path doesn't need correct data — only the
    # number of bytes per DMA matters — so bypass DataExecutor for the
    # chunked kernel. Composite/naive-async loads use the full (M, K)
    # shape and remain data-correct.
    if bench_name in ("matmul-async-chunked", "matmul-async-chunked-db"):
        GraphEngine._flush_data_phase = lambda self: None

    topo = resolve_topology(str(TOPO_PATH))
    bench = resolve_bench(bench_name).run
    device = resolve_device(None)
    result = run_bench(
        topology=topo, bench_fn=bench, device=device,
        engine_factory=lambda t, d: GraphEngine(
            getattr(t, "topology_obj", t), enable_data=True,
        ),
    )
    if not result.completion.ok:
        raise RuntimeError(f"{bench_name} failed at {M}x{K}x{N}: {result.completion}")
    log = result.engine.op_log
    pe_window = _pe_window_ns(log)
    if bench_name == "matmul-composite":
        engine_window = _composite_window_ns(log)
    elif bench_name in ("matmul-async-chunked", "matmul-async-chunked-db"):
        # First N_chunks pe_dma records are A pre-stage; engine window
        # starts at the (N_chunks+1)-th pe_dma record (= first B-chunk
        # load). Fall back to the naive analog for K <= TILE_K (kernel
        # collapses to single load+dot+store).
        TILE_K = 64
        n_chunks = max(K // TILE_K, 1)
        if n_chunks <= 1:
            engine_window = _async_engine_window_ns(log)
        else:
            dma = sorted(
                (r for r in log if r.component_id.endswith(".pe_dma")),
                key=lambda r: r.t_start,
            )
            pe = _pe_records(log)
            if len(dma) > n_chunks and pe:
                engine_window = max(r.t_end for r in pe) - dma[n_chunks].t_start
            else:
                engine_window = _async_engine_window_ns(log)
    else:
        engine_window = _async_engine_window_ns(log)
    flops = 2 * M * K * N
    # flops / ns * 1e-3 = TFLOP/s  (since 1 flop/ns = 1 GFLOP/s)
    return {
        "M": M, "K": K, "N": N,
        "bench": bench_name, "variant": variant,
        "pe_window_ns": pe_window,
        "engine_window_ns": engine_window,
        "flops": flops,
        "tflops": (flops / engine_window / 1000.0) if engine_window > 0 else 0.0,
        "n_records": len(log),
    }


def collect() -> list[dict]:
    rows: list[dict] = []
    for M, K, N in SHAPES:
        print(f"  shape M={M:4d} K={K:5d} N={N:4d} ...", flush=True)
        comp = _run_one("matmul-composite", "load_ref", M, K, N)
        asyn = _run_one("matmul-async", None, M, K, N)
        chnk = _run_one("matmul-async-chunked", None, M, K, N)
        chnkdb = _run_one("matmul-async-chunked-db", None, M, K, N)
        print(f"    composite:    {comp['engine_window_ns']:8.1f} ns  "
              f"{comp['tflops']:6.3f} TFLOPS")
        print(f"    async naive:  {asyn['engine_window_ns']:8.1f} ns  "
              f"{asyn['tflops']:6.3f} TFLOPS")
        print(f"    chunked all:  {chnk['engine_window_ns']:8.1f} ns  "
              f"{chnk['tflops']:6.3f} TFLOPS")
        print(f"    chunked db=2: {chnkdb['engine_window_ns']:8.1f} ns  "
              f"{chnkdb['tflops']:6.3f} TFLOPS")
        rows.append(comp)
        rows.append(asyn)
        rows.append(chnk)
        rows.append(chnkdb)
    return rows


def plot(rows: list[dict]) -> None:
    import matplotlib.pyplot as plt
    import numpy as np

    shape_keys = []
    for M, K, N in SHAPES:
        shape_keys.append((M, K, N))

    by_key: dict[tuple[int, int, int], dict[str, dict]] = {}
    for r in rows:
        key = (r["M"], r["K"], r["N"])
        by_key.setdefault(key, {})[r["bench"]] = r

    labels = [f"M={M}\nK={K}\nN={N}" for (M, K, N) in shape_keys]
    comp_tflops = [by_key[k]["matmul-composite"]["tflops"] for k in shape_keys]
    asyn_tflops = [by_key[k]["matmul-async"]["tflops"] for k in shape_keys]
    chnk_tflops = [by_key[k]["matmul-async-chunked"]["tflops"] for k in shape_keys]
    chnkdb_tflops = [by_key[k]["matmul-async-chunked-db"]["tflops"]
                     for k in shape_keys]

    x = np.arange(len(shape_keys))
    width = 0.20

    fig, ax = plt.subplots(figsize=(14, 5.5))
    b_c = ax.bar(x - 1.5*width, comp_tflops, width,
                 label="Composite (load_ref)", color="#10b981")
    b_a = ax.bar(x - 0.5*width, asyn_tflops, width,
                 label="Async naive (tl.load full + tl.dot)",
                 color="#f59e0b")
    b_db = ax.bar(x + 0.5*width, chnkdb_tflops, width,
                  label="Async chunked-prefetch, depth=2 "
                        "(TCM-bounded)", color="#a855f7")
    b_k = ax.bar(x + 1.5*width, chnk_tflops, width,
                 label="Async chunked-prefetch, depth=$\\infty$ "
                       "(all B-tiles queued up front)", color="#3b82f6")

    ax.axhline(8.0, linestyle="--", color="#94a3b8", linewidth=0.8,
               label="Per-PE GEMM peak (8 TFLOP/s)")
    ax.set_ylabel("Per-PE achieved TFLOP/s")
    ax.set_title("Composite vs async-loading GEMM — per-PE throughput "
                 "(engine pipeline window, A pre-stage excluded)")
    ax.set_xticks(x)
    ax.set_xticklabels(labels, fontsize=8)
    ax.grid(axis="y", alpha=0.25)
    ax.legend(loc="upper left", fontsize=9)

    for bar in (*b_c, *b_a, *b_db, *b_k):
        h = bar.get_height()
        ax.annotate(f"{h:.2f}", xy=(bar.get_x() + bar.get_width()/2, h),
                    xytext=(0, 2), textcoords="offset points",
                    ha="center", va="bottom", fontsize=7, color="#475569")

    fig.tight_layout()
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    fig.savefig(OUT_PNG, dpi=150)
    print(f"wrote {OUT_PNG.relative_to(REPO)}")
    plt.close(fig)


def main() -> int:
    rows = collect()
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    OUT_JSON.write_text(json.dumps(rows, indent=2))
    print(f"wrote {OUT_JSON.relative_to(REPO)}")
    plot(rows)
    return 0


if __name__ == "__main__":
    sys.exit(main())