#!/usr/bin/env python3
"""SCRATCH EXPERIMENT (not production; do not commit).

Question: does charging the *primitive* decode kernel for per-HW-tile
(16x16x16) CPU dispatch flip the "composite gives no decode-latency
benefit" conclusion?

We monkeypatch TLContext.dot so that, in the primitive kernel, every
tl.dot whose (M,K,N) exceeds the HW GEMM tile (mac_m/mac_k/mac_n) is
split by the CPU into ceil(M/mac_m)*ceil(K/mac_k)*ceil(N/mac_n)
HW-tile-sized GemmCmds.  Each tile GemmCmd is emitted through the normal
_emit() path, so it (a) charges PE_CPU dispatch overhead via
_charge_dispatch (PeCpuOverheadCmd), and (b) blocks like a normal
single-op GemmCmd on PE_GEMM at the cycle-accurate ceil-product latency.

We also inject mac_m/mac_k/mac_n into every pe_gemm topology node so BOTH
the primitive-tiled and the composite variants run on the *same*
cycle-accurate engine (fair comparison).  Composite is left untouched:
the CPU emits ONE CompositeCmd, and PE_SCHEDULER tiles internally (no
per-HW-tile CPU dispatch).

Data correctness:
  Inputs are ctx.zeros (q/k/v), so every matmul result is zeros and the
  DataExecutor replay is trivial.  To keep replay numerically correct
  regardless, exactly ONE emitted tile per dot carries the *real* full
  operands+output handles (so the DataExecutor computes the true (M,N)
  result via the recorded handle shapes), while its timing fields
  (m,k,n) are the HW-tile size so the engine charges exactly one tile of
  cycle time.  The remaining n_tiles-1 emitted tiles are timing-only
  GemmCmds (16x16x16) writing to throwaway scratch.  Net: n_tiles tiles
  of engine time + n_tiles dispatch charges, and a correct final output.

Engine mode: enable_data=True (same as the production sweep's
_engine_latency_ns), op_log end-to-end latency.
"""
from __future__ import annotations

import json
from math import ceil
from pathlib import Path

# --- HW GEMM tile under test -------------------------------------------------
MAC_M = 16
MAC_K = 16
MAC_N = 16
# Legacy alt to also try: (8, 16, 32)

S_KV_LATENCY = (8192, 32_768, 65_536, 131_072)

ROOT = Path(__file__).resolve().parent
SWEEP_JSON = (
    ROOT / "src" / "kernbench" / "benches" / "1H_milestone_output"
    / "gqa" / "long_ctx" / "sweep_decode_composite.json"
)


# --- mac-dim topology override -----------------------------------------------
def _topo_with_mac(mac_m: int, mac_k: int, mac_n: int):
    """Compiled topology with mac dims injected into every pe_gemm node."""
    from kernbench.topology.builder import resolve_topology

    handle = resolve_topology("topology.yaml")
    g = handle.topology_obj
    n = 0
    for node in g.nodes.values():
        if node.kind == "pe_gemm":
            node.attrs["mac_m"] = mac_m
            node.attrs["mac_k"] = mac_k
            node.attrs["mac_n"] = mac_n
            n += 1
    print(f"    injected mac=({mac_m},{mac_k},{mac_n}) into {n} pe_gemm nodes")
    return handle


# --- tiling monkeypatch for TLContext.dot ------------------------------------
def _make_tiled_dot(orig_dot, mac_m: int, mac_k: int, mac_n: int):
    from kernbench.common.pe_commands import GemmCmd

    def tiled_dot(self, a, b):
        if len(a.shape) < 2 or len(b.shape) < 2:
            return orig_dot(self, a, b)
        m, k = a.shape[-2], a.shape[-1]
        k2, n = b.shape[-2], b.shape[-1]
        if k != k2:
            raise ValueError(f"dot shape mismatch: a.K={k} != b.K={k2}")

        n_tiles = ceil(m / mac_m) * ceil(k / mac_k) * ceil(n / mac_n)

        out_shape = (*a.shape[:-2], m, n)
        out = self._make_compute_out(shape=out_shape, dtype=a.dtype)
        self._await_pending(a, b)

        if n_tiles <= 1:
            self._emit(GemmCmd(a=a, b=b, out=out, m=m, k=k, n=n))
            return out

        # One real-data tile: full handles (so DataExecutor computes the
        # true result), but timing fields = HW tile (one tile of cycles).
        self._emit(GemmCmd(a=a, b=b, out=out, m=mac_m, k=mac_k, n=mac_n))
        # Remaining timing-only tiles: throwaway scratch, 16x16x16.
        scratch = self._make_compute_out(shape=(mac_m, mac_n), dtype=a.dtype)
        for _ in range(n_tiles - 1):
            self._emit(GemmCmd(a=a, b=b, out=scratch,
                               m=mac_m, k=mac_k, n=mac_n))
        return out

    return tiled_dot


# --- latency runner (replicates sweep's _engine_latency_ns) ------------------
def _engine_latency_ns(variant: str, S_kv: int, topo) -> float:
    from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (  # noqa: E501
        _end_to_end_ns, _run_panel_fn,
    )
    from kernbench.runtime_api.bench_runner import run_bench
    from kernbench.runtime_api.types import resolve_device
    from kernbench.sim_engine.engine import GraphEngine

    result = run_bench(
        topology=topo, bench_fn=_run_panel_fn(variant, S_kv),
        device=resolve_device(None),
        engine_factory=lambda t, d: GraphEngine(
            getattr(t, "topology_obj", t), enable_data=True,
        ),
    )
    if not result.completion.ok:
        raise RuntimeError(
            f"{variant}@{S_kv} failed: {result.completion}"
        )
    return _end_to_end_ns(result.engine.op_log)


def _emit_dispatch(variant: str, S_kv: int) -> int:
    """PE_CPU command count at the center rank (cube 6, pe 0)."""
    from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (  # noqa: E501
        _emit_dispatch as prod_emit,
    )
    return prod_emit(variant, S_kv)[0]


def main() -> None:
    import kernbench.triton_emu.tl_context as tlc

    # Baseline (A): primitive UNTILED latencies from the production sweep
    # (mac=0 / TFLOPS model). Read straight off the committed sweep JSON.
    sweep = json.loads(SWEEP_JSON.read_text())
    base_A = {}
    for r in sweep["rows"]:
        if r["variant"] == "primitive" and r["latency_ns"] is not None:
            base_A[r["S_kv"]] = r["latency_ns"]

    print(f"== mac tile = ({MAC_M},{MAC_K},{MAC_N}) ==")

    # --- command-count sanity (emit-time, mac-independent) ---------------
    orig_dot = tlc.TLContext.dot
    print("\n[dispatch counts @ S_kv=131072]")
    n_prim_untiled = _emit_dispatch("primitive", 131072)
    tlc.TLContext.dot = _make_tiled_dot(orig_dot, MAC_M, MAC_K, MAC_N)
    try:
        n_prim_tiled = _emit_dispatch("primitive", 131072)
        n_comp = None
    finally:
        tlc.TLContext.dot = orig_dot
    n_comp = _emit_dispatch("composite", 131072)
    print(f"  primitive UNTILED PE_CPU cmds : {n_prim_untiled}")
    print(f"  primitive TILED   PE_CPU cmds : {n_prim_tiled}  "
          f"(x{n_prim_tiled / max(n_prim_untiled,1):.0f})")
    print(f"  composite         PE_CPU cmds : {n_comp}")

    # --- latency sweep ----------------------------------------------------
    rows = []
    topo = _topo_with_mac(MAC_M, MAC_K, MAC_N)

    for S_kv in S_KV_LATENCY:
        A = base_A.get(S_kv)

        # (C) composite on the mac engine (untouched dot path)
        C = _engine_latency_ns("composite", S_kv, topo)

        # (B) primitive TILED on the mac engine
        tlc.TLContext.dot = _make_tiled_dot(orig_dot, MAC_M, MAC_K, MAC_N)
        try:
            B = _engine_latency_ns("primitive", S_kv, topo)
        finally:
            tlc.TLContext.dot = orig_dot

        gap_pct = (B - C) / C * 100.0 if C else float("nan")
        rows.append((S_kv, A, B, C, gap_pct))
        print(f"  S_kv={S_kv:>7}: A(untiled)={A!s:>12}  "
              f"B(tiled)={B:12.2f}  C(comp)={C:12.2f}  (B-C)/C={gap_pct:+6.1f}%")

    # --- final table ------------------------------------------------------
    print("\n==================== RESULT TABLE ====================")
    print(f"{'S_kv':>8} | {'A untiled(ns)':>14} | {'B tiled(ns)':>14} | "
          f"{'C comp(ns)':>14} | {'(B-C)/C':>9}")
    print("-" * 72)
    for S_kv, A, B, C, gap in rows:
        a_s = f"{A:.2f}" if A is not None else "n/a"
        print(f"{S_kv:>8} | {a_s:>14} | {B:>14.2f} | {C:>14.2f} | "
              f"{gap:>+8.1f}%")


if __name__ == "__main__":
    main()