"""A/B the vendored FlagGems 5.0.2 kernel against the TileLang kernel, 7 shapes.

Same cases as scripts/sparse_mla_port/ab_vs_tilelang.py, but importing the in-tree copy instead of
the downloaded one, so this is what actually runs in the engine.
"""

import math
import sys

import torch

from deepinfer.h20_ops.sparse_mla_sm90 import sparse_mla_fwd_attn
from deepinfer.h20_ops.vendor.flashmla_sparse import flash_mla_sparse_fwd

DEV = "cuda"
HQ, HKV, D, TOPK = 64, 1, 512, 512
FAILED = []


def tl_call(q3: torch.Tensor, kv3: torch.Tensor, idx3: torch.Tensor, scale: float) -> tuple[torch.Tensor, torch.Tensor]:
    s_q, s_kv = q3.shape[0], kv3.shape[0]
    out, lse = sparse_mla_fwd_attn(
        q3.view(1, s_q, HQ, D),
        kv3.view(1, s_kv, 1, D),
        idx3.view(1, s_q, 1, -1),
        sm_scale=scale,
        is_causal=False,
    )
    return out.view(s_q, HQ, D), lse.view(s_q, HQ)


def compare(tag: str, a: torch.Tensor, b: torch.Tensor) -> None:
    af, bf = a.float(), b.float()
    if af.shape != bf.shape:
        print(f"  FAIL {tag}: shape {tuple(af.shape)} vs {tuple(bf.shape)}")
        FAILED.append(tag)
        return
    diff = (af - bf).abs()
    scale = af.abs().max().item()
    tol = 2.0**-8 * max(scale, 1e-9) * 2
    ok = diff.max().item() <= tol and bool(torch.isfinite(bf).all())
    print(f"  {'ok  ' if ok else 'FAIL'} {tag}: max|ref| {scale:.4f}  max diff {diff.max().item():.3e}  tol {tol:.3e}")
    if not ok:
        FAILED.append(tag)


def case(tag: str, SQ: int, SKV: int, seed: int, distinct: bool = True) -> None:
    torch.manual_seed(seed)
    q = (torch.randn((SQ, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    kv = (torch.randn((SKV, HKV, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    if distinct and SKV > TOPK:
        idx = torch.stack([torch.randperm(SKV, device=DEV)[:TOPK] for _ in range(SQ)])
        idx = idx.view(SQ, 1, TOPK).to(torch.int32)
    else:
        idx = torch.randint(0, SKV, (SQ, 1, TOPK), device=DEV, dtype=torch.int32)
    idx = idx.contiguous()

    print(f"\n[{tag}] SQ={SQ} SKV={SKV} distinct={distinct}")
    scale = D**-0.5
    tl_out, tl_lse = tl_call(q, kv, idx, scale)
    torch.cuda.synchronize()
    fg_out, fg_max, fg_lse = flash_mla_sparse_fwd(q, kv, idx, sm_scale=scale, d_v=D)
    torch.cuda.synchronize()
    compare("out", tl_out, fg_out)
    # lse units differ: TileLang returns log2, FlagGems natural log.
    compare("lse (after log2->ln)", tl_lse.float() * math.log(2), fg_lse)


CASES = [
    ("prefill, distinct top-k", 64, 2048, 0, True),
    ("prefill, duplicate top-k", 32, 512, 1, False),
    ("decode-like (SQ=1)", 1, 4096, 2, True),
    ("odd SQ", 7, 1024, 3, True),
    ("short KV < topk", 16, 128, 4, False),
    ("larger prefill", 256, 4096, 5, True),
    ("equal SQ==SKV", 512, 512, 6, False),
]

if __name__ == "__main__":
    for c in CASES:
        case(*c)
    print()
    if FAILED:
        print(f"MISMATCH {len(FAILED)}: {FAILED}")
        sys.exit(1)
    print(f"ALL_MATCH ({len(CASES)} cases)")