"""Head-to-head against upstream's TileLang Mamba-3 MIMO forward.

The two kernels take different dtype contracts at reduced precision. Upstream
keeps the biases, the rank weights and D in float32 and casts trap to the value
dtype; this kernel does the reverse, carrying the biases and rank weights in the
value dtype and the whole schedule in float32. So one float32 master set is
built and each side is handed the cast it requires, which keeps the underlying
values identical.

In float32 the two contracts coincide, so that mode doubles as a correctness
anchor. Both paths run in one process on one device, and TileLang autotunes per
shape, so everything is warmed before it is timed.

    PYTHONPATH=~/k/mamba-upstream ~/m3bench/bin/python bench_vs_upstream.py
"""

import argparse
import sys
import time
from pathlib import Path

import torch
import torch.nn.functional as F

M3 = Path.home() / "k" / "mamba3"
UP = Path.home() / "k" / "mamba-upstream"
sys.path.insert(0, str(M3))
sys.path.insert(0, str(UP))

import load_local as mamba3  # noqa: E402
from mamba_ssm.ops.tilelang.mamba3.mamba3_mimo import mamba3_mimo  # noqa: E402

# ours: value dtype on the activations, the biases and the rank weights
OURS_CAST = ["q", "k", "v", "z", "q_bias", "k_bias", "mimo_v", "mimo_o", "mimo_z", "D"]
# upstream: value dtype on the activations and trap only
THEIRS_CAST = ["q", "k", "v", "z", "trap"]


def master(B, S, H, G, P, N, R, Na, dev, seed):
    """One float32 set of values; each side casts what it needs."""
    torch.manual_seed(seed)
    f = lambda *s: torch.randn(*s, device=dev)
    dt = F.softplus(-3.0 + f(B, H, S))
    return {
        "q": f(B, S, R, G, N), "k": f(B, S, R, G, N),
        "v": f(B, S, H, P), "z": f(B, S, H, P),
        "q_bias": f(H, R, N), "k_bias": f(H, R, N),
        "mimo_v": torch.rand(H, R, P, device=dev) / R,
        "mimo_o": torch.rand(H, R, P, device=dev) / R,
        "mimo_z": torch.rand(H, R, P, device=dev) / R,
        "D": f(H), "angles": torch.rand(B, S, H, Na, device=dev),
        "dt": dt, "adt": -F.softplus(f(B, H, S)).clamp(max=-1e-4) * dt,
        "trap": torch.rand(B, H, S, device=dev) * 0.5,
    }


def cast(c, names, dtype):
    if dtype is torch.float32:
        return dict(c)
    return {k: (t.to(dtype).contiguous() if k in names else t) for k, t in c.items()}


def to_split_half(c, N):
    """Reindex q, k and their biases from pairwise rotary layout to split-half.

    This kernel rotates the pair (2i, 2i+1); the TileLang kernel rotates
    (n, N/2 + n). The two are a permutation of the head dimension, and it is
    applied to q and k alike, so every inner product is preserved and the two
    compute the same operator on the same values.
    """
    idx = torch.cat([torch.arange(0, N, 2), torch.arange(1, N, 2)]).to(c["q"].device)
    out = dict(c)
    for k in ("q", "k", "q_bias", "k_bias"):
        out[k] = c[k].index_select(-1, idx).contiguous()
    return out


def ours(c, C):
    return mamba3.forward(
        c["q"], c["k"], c["v"], c["q_bias"], c["k_bias"], c["mimo_v"], c["mimo_o"],
        c["angles"], c["adt"], c["dt"], c["trap"], z=c["z"], mimo_z=c["mimo_z"],
        D=c["D"], chunk_size=C)


def theirs(c, C, dtype):
    return mamba3_mimo(
        c["q"], c["k"], c["v"], c["adt"], c["dt"], c["trap"],
        c["q_bias"], c["k_bias"], c["mimo_v"], c["mimo_z"], c["mimo_o"],
        c["angles"], c["D"], c["z"],
        chunk_size=C, rotary_dim_divisor=2, dtype=dtype)


def timeit(fn, iters, warmup=5):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(iters):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / iters * 1e3


def rel(a, b):
    a, b = a.float(), b.float()
    return float((a - b).abs().max() / b.abs().max().clamp_min(1e-5))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float32"])
    ap.add_argument("--iters", type=int, default=20)
    args = ap.parse_args()
    dtype = getattr(torch, args.dtype)
    dev = "cuda"

    cap = "".join(map(str, torch.cuda.get_device_capability()))
    print(f"torch {torch.__version__}  {torch.cuda.get_device_name(0)}  "
          f"sm_{cap}  dtype={args.dtype}\n")
    print(f"{'S':>6} {'H':>4} {'P':>4} {'N':>4} {'R':>3} {'C':>4} "
          f"{'ours ms':>10} {'tilelang ms':>12} {'ratio':>8}  {'max rel':>9}")

    B, G = 1, 1
    grid = [(512, 8, 64, 128, 4, 16), (1024, 8, 64, 128, 4, 16),
            (2048, 8, 64, 128, 4, 16), (4096, 8, 64, 128, 4, 16),
            (2048, 32, 64, 128, 4, 16), (2048, 8, 64, 128, 1, 64)]

    for S, H, P, N, R, C in grid:
        Na = N // 2
        m = master(B, S, H, G, P, N, R, Na, dev, seed=0)
        co = cast(m, OURS_CAST, dtype)
        ct = cast(to_split_half(m, N), THEIRS_CAST, dtype)
        try:
            yo, yt = ours(co, C), theirs(ct, C, dtype)
            r = rel(yo, yt)
            to = timeit(lambda: ours(co, C), args.iters)
            tt = timeit(lambda: theirs(ct, C, dtype), args.iters)
            print(f"{S:6d} {H:4d} {P:4d} {N:4d} {R:3d} {C:4d} "
                  f"{to:10.3f} {tt:12.3f} {tt / to:7.2f}x  {r:9.2e}")
        except Exception as exc:
            msg = f"{type(exc).__name__}: {exc}".replace("\n", " ")
            print(f"{S:6d} {H:4d} {P:4d} {N:4d} {R:3d} {C:4d}   {msg[:88]}")
        finally:
            del m, co, ct
            torch.cuda.empty_cache()

    return 0


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