"""05_benchmark.py -- reproduce the perf numbers in the README.

Runs the SISO and MIMO forward + decode benchmarks at target dims. Reports
warm median latency + p95, and a summary table matching the README.

Run:
  python 05_benchmark.py [--full]

  --full: also runs seqlen=512 and seqlen=1024 (adds ~5 minutes)

Note: MIMO does not use torch.compile due to a known regression -- see the
project's ticket 61. SISO uses torch.compile at batch=1 for optimal single-request
latency.
"""

import argparse
import statistics
import time
import torch
from _loader import load_kernel


def _sync():
    try:
        torch.neuron.synchronize()
    except Exception:
        pass


def bench_forward(mixer, batch, seqlen, warmup=3, iters=10):
    u = torch.randn(batch, seqlen, 1024, device="neuron")
    for _ in range(warmup):
        y, _ = mixer(u)
        _sync()
    times = []
    for _ in range(iters):
        t0 = time.perf_counter()
        y, _ = mixer(u)
        _sync()
        times.append((time.perf_counter() - t0) * 1000)
    return statistics.median(times), sorted(times)[int(len(times) * 0.95)]


def bench_decode(mixer, batch, prefill_len, n_steps=32, warmup=3):
    u_pre = torch.randn(batch, prefill_len, 1024, device="neuron")
    _, cache = mixer(u_pre)
    _sync()
    tokens = [torch.randn(batch, 1, 1024, device="neuron") for _ in range(warmup + n_steps)]
    for tok in tokens[:warmup]:
        y, cache = mixer.step(tok, cache)
        _sync()
    times = []
    for tok in tokens[warmup:]:
        t0 = time.perf_counter()
        y, cache = mixer.step(tok, cache)
        _sync()
        times.append((time.perf_counter() - t0) * 1000)
    return statistics.median(times), sorted(times)[int(len(times) * 0.95)]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--full", action="store_true", help="Include seqlen 512 and 1024")
    args = ap.parse_args()

    print("=" * 78)
    print("mamba3-neuron-kernels benchmark  (target dims)")
    print("=" * 78)
    print("Setup: d_model=1024, d_state=128, headdim=64, nheads=32, batch=1")
    print("Env:  TORCH_NEURONX_ENABLE_CONCATENATION=1 recommended")

    mamba3 = load_kernel()

    seqlens = [64, 256, 512, 1024] if args.full else [64, 256]

    # ---- SISO Forward ----
    print("\n### SISO forward (NKI SSD + torch.compile) ###")
    print(f"{'seqlen':>7s} {'median (ms)':>12s} {'p95 (ms)':>10s}")
    print("-" * 32)
    for seqlen in seqlens:
        torch.manual_seed(0)
        m = mamba3.NeuronMamba3Mixer(
            d_model=1024, d_state=128, headdim=64, chunk_size=64, mimo_rank=1,
            use_nki_ssd=True,
        ).to("neuron")
        m_compiled = torch.compile(m, backend="neuron")
        try:
            med, p95 = bench_forward(m_compiled, 1, seqlen)
            print(f"{seqlen:>7d} {med:>10.2f}   {p95:>8.2f}")
        except Exception as e:
            print(f"{seqlen:>7d} FAILED: {type(e).__name__}: {str(e)[:40]}")
        del m, m_compiled

    # ---- MIMO Forward ----
    print("\n### MIMO forward (NKI SSD, eager -- torch.compile disabled per ticket 61) ###")
    print(f"{'seqlen':>7s} {'median (ms)':>12s} {'p95 (ms)':>10s}")
    print("-" * 32)
    for seqlen in seqlens:
        torch.manual_seed(0)
        m = mamba3.NeuronMamba3Mixer(
            d_model=1024, d_state=128, headdim=64, chunk_size=16, mimo_rank=4,
            use_nki_ssd=True,
        ).to("neuron")
        try:
            med, p95 = bench_forward(m, 1, seqlen)
            print(f"{seqlen:>7d} {med:>10.2f}   {p95:>8.2f}")
        except Exception as e:
            print(f"{seqlen:>7d} FAILED: {type(e).__name__}: {str(e)[:40]}")
        del m

    # ---- Decode step ----
    print("\n### Decode step (single-token, prefill_len=64) ###")
    print(f"{'mode':>6s} {'median (ms)':>12s} {'p95 (ms)':>10s}")
    print("-" * 32)
    for mode in ["SISO", "MIMO"]:
        torch.manual_seed(0)
        if mode == "SISO":
            m = mamba3.NeuronMamba3Mixer(
                d_model=1024, d_state=128, headdim=64, chunk_size=64, mimo_rank=1,
                use_nki_ssd=True,
            ).to("neuron")
        else:
            m = mamba3.NeuronMamba3Mixer(
                d_model=1024, d_state=128, headdim=64, chunk_size=16, mimo_rank=4,
                use_nki_ssd=True,
            ).to("neuron")
        med, p95 = bench_decode(m, 1, prefill_len=64, n_steps=32)
        print(f"{mode:>6s} {med:>10.2f}   {p95:>8.2f}")
        del m

    print("\n" + "=" * 78)
    print("Done.")


if __name__ == "__main__":
    main()