"""
Sweep + plot harness for the NVFP4 Dual-GEMM+SiLU benchmark.

Times all three methods (CuTe fused kernel, vLLM cutlass_scaled_fp4_mm, flashinfer
mm_fp4) and writes three figures as PNGs:

  1. bench_named_shapes.png  -- grouped TFLOPS bars for the 3 named Laguna shapes
  2. bench_sweep_sx2.png     -- TFLOPS vs M (256..4096) for SX2 (N=512,  K=2048)
  3. bench_sweep_m1.png      -- TFLOPS vs M (256..4096) for M1  (N=4096, K=16384)

The CuTe point is only plotted when the kernel actually matches the PyTorch
reference for that shape; otherwise it is skipped (so a fall-back / wrong result
never shows up as a bogus number). vLLM / flashinfer points are dropped if the
library is unavailable or errors.

Run on the B200 remote:   python bench_sweep_plot.py
NOTE: do not rename this file to flashinfer.py -- it would shadow the package.
"""
import torch
import matplotlib
matplotlib.use("Agg")  # headless
import matplotlib.pyplot as plt

from bench_flashinfer_2 import (
    generate_input,
    ref_kernel,
    _benchmark_fn,
    _prepare_vllm_dual_gemm_silu,
    _prepare_flashinfer_dual_gemm_silu,
)


def _cute_quality(data, out, rtol=1e-3, atol=1e-3, max_bad_frac=1e-4):
    """
    Classify the CuTe result against the PyTorch reference for *benchmarking*.

    Returns (accept, note). We gate on the FRACTION of mismatched elements, not
    their magnitude:
      * a handful of elements differ -> non-determinism / SiLU-kink / large-K
        accumulation rounding. Harmless for a FLOPS measurement -> accept, plot.
      * a large fraction is wrong     -> real kernel bug (e.g. the earlier
        unwritten-tiles case was ~94% of elements) -> reject, skip the point.
    """
    ref = ref_kernel(data).float()
    diff = (out.float() - ref).abs()
    bad = diff > (atol + rtol * ref.abs())
    n_bad = int(bad.sum().item())
    total = ref.numel()
    if n_bad == 0:
        return True, "exact"
    frac = n_bad / total
    max_abs = float(diff.max().item())
    if frac < max_bad_frac and max_abs == max_abs and max_abs != float("inf"):
        return True, (f"{n_bad}/{total} elems differ (max|Δ|={max_abs:.1e}); "
                      f"ignored for FLOPS")
    return False, f"{n_bad}/{total} mismatched (max|Δ|={max_abs:.1e})"

METHODS = ["CuTe DualGEMM", "vLLM", "Flashinfer"]
COLORS = {"CuTe DualGEMM": "tab:blue", "vLLM": "tab:orange", "Flashinfer": "tab:green"}

# SX2 = (N=512, K=2048), M1 = (N=4096, K=16384)
SX2_NK = (512, 2048)
M1_NK = (512 * 8, 2048 * 8)
SWEEP_M = [256, 512, 1024, 2048, 4096]

NAMED = [
    ("SX2 prefill\n4096x512x2048", 4096, SX2_NK[0], SX2_NK[1]),
    ("M1 prefill\n4096x4096x16384", 4096, M1_NK[0], M1_NK[1]),
    ("SX2 decode\n256x512x2048", 256, SX2_NK[0], SX2_NK[1]),
]


def measure_one(m, n, k, l=1, seed=42, warmup=20, iters=100):
    """Return {method: TFLOPS or None} for a single shape (None = skipped/failed)."""
    from laguna_dual_gemm import custom_kernel

    data = generate_input(m, n, k, l, seed=seed)
    a, b1, b2, sfa, sfb1, sfb2, sfa_p, sfb1_p, sfb2_p, c = data
    flops = 2 * 2 * m * n * k * l  # two GEMMs, 2*M*N*K each
    res = {meth: None for meth in METHODS}

    # --- CuTe (counted if it does the real computation; fp noise tolerated) ---
    try:
        custom_out = custom_kernel(data).clone()
        accept, note = _cute_quality(data, custom_out)
        if accept:
            ms = _benchmark_fn(lambda: custom_kernel(data), warmup=warmup, iters=iters)
            res["CuTe DualGEMM"] = flops / (ms * 1e-3) / 1e12
            if note != "exact":
                print(f"    [CuTe]       {m}x{n}x{k}: accepted ({note})")
        else:
            print(f"    [CuTe]       skip {m}x{n}x{k}: {note}")
    except Exception as e:
        print(f"    [CuTe]       skip {m}x{n}x{k}: {e}")

    # --- vLLM ---
    try:
        run = _prepare_vllm_dual_gemm_silu(a, b1, b2, sfa, sfb1, sfb2, m, n, k, l)
        if run is None:
            print(f"    [vLLM]       skip {m}x{n}x{k}: unavailable")
        else:
            run(); torch.cuda.synchronize()
            ms = _benchmark_fn(run, warmup=warmup, iters=iters)
            res["vLLM"] = flops / (ms * 1e-3) / 1e12
    except Exception as e:
        print(f"    [vLLM]       skip {m}x{n}x{k}: {e}")

    # --- Flashinfer ---
    try:
        run = _prepare_flashinfer_dual_gemm_silu(a, b1, b2, sfa, sfb1, sfb2, m, n, k, l)
        if run is None:
            print(f"    [Flashinfer] skip {m}x{n}x{k}: unavailable")
        else:
            run(); torch.cuda.synchronize()
            ms = _benchmark_fn(run, warmup=warmup, iters=iters)
            res["Flashinfer"] = flops / (ms * 1e-3) / 1e12
    except Exception as e:
        print(f"    [Flashinfer] skip {m}x{n}x{k}: {e}")

    line = "  ".join(
        f"{meth}={res[meth]:.1f}" if res[meth] is not None else f"{meth}=--"
        for meth in METHODS
    )
    print(f"  M={m:>5} N={n} K={k}: {line}")
    return res


def sweep(n, k, ms):
    """Return {method: [tflops per m]} (None where skipped)."""
    out = {meth: [] for meth in METHODS}
    for m in ms:
        r = measure_one(m, n, k)
        for meth in METHODS:
            out[meth].append(r[meth])
    return out


def plot_sweep(ms, data, title, fname):
    plt.figure(figsize=(7, 5))
    for meth in METHODS:
        xs = [m for m, y in zip(ms, data[meth]) if y is not None]
        ys = [y for y in data[meth] if y is not None]
        if ys:
            plt.plot(xs, ys, marker="o", linewidth=2, label=meth, color=COLORS[meth])
    plt.xscale("log", base=2)
    plt.xticks(ms, [str(m) for m in ms])
    plt.xlabel("M (tokens)")
    plt.ylabel("TFLOPS")
    plt.title(f"NVFP4 Dual-GEMM+SiLU  --  {title}")
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.savefig(fname, dpi=130)
    plt.close()
    print(f"wrote {fname}")


def plot_named(named_results, fname):
    import numpy as np

    labels = [lbl for lbl, *_ in NAMED]
    x = np.arange(len(labels))
    width = 0.25
    plt.figure(figsize=(9, 5))
    for i, meth in enumerate(METHODS):
        heights = [(r[meth] if r[meth] is not None else 0.0) for r in named_results]
        bars = plt.bar(x + (i - 1) * width, heights, width, label=meth, color=COLORS[meth])
        for b, r in zip(bars, named_results):
            if r[meth] is None:
                plt.text(b.get_x() + b.get_width() / 2, 0, "N/A",
                         ha="center", va="bottom", fontsize=8, rotation=90)
    plt.xticks(x, labels)
    plt.ylabel("TFLOPS")
    plt.title("NVFP4 Dual-GEMM+SiLU  --  named Laguna shapes")
    plt.grid(True, axis="y", alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.savefig(fname, dpi=130)
    plt.close()
    print(f"wrote {fname}")


if __name__ == "__main__":
    print("=== TFLOPS-vs-M sweep: SX2 (N=512, K=2048) ===")
    sx2 = sweep(SX2_NK[0], SX2_NK[1], SWEEP_M)

    print("=== TFLOPS-vs-M sweep: M1 (N=4096, K=16384) ===")
    m1 = sweep(M1_NK[0], M1_NK[1], SWEEP_M)

    print("=== named shapes ===")
    named_results = [measure_one(m, n, k) for _, m, n, k in NAMED]

    plot_sweep(SWEEP_M, sx2, "SX2 (N=512, K=2048)", "bench_sweep_sx2.png")
    plot_sweep(SWEEP_M, m1, "M1 (N=4096, K=16384)", "bench_sweep_m1.png")
    plot_named(named_results, "bench_named_shapes.png")
    print("\nDone. 3 figures written: "
          "bench_sweep_sx2.png, bench_sweep_m1.png, bench_named_shapes.png")