"""Unit tests for the Triton sparse-MLA adapter, against the TileLang kernel as reference.

The adapter's job is not the attention maths -- that is the vendored kernel's, and
`ab_vendored.py` covers it. The job is the *contract*: return arity, lse units, dtype, and the two
calling conventions `attn_sm90` uses. Those are what this file tests, because each of them fails
silently if wrong:

* returning natural-log lse makes `_apply_sink` scale outputs by exp(lse*(1-ln2)) -- plausible
  output, wrong numbers;
* returning (out, max_logits, lse) instead of (out, lse) makes every `out, lse = ...` unpack the
  wrong tensor;
* mismapping the per-query convention attends the wrong query's window.

Run: PYTHONPATH=. python scripts/sparse_mla_port/test_sparse_mla_triton.py
"""

from __future__ import annotations

import math
import sys

import torch

from deepinfer.h20_ops import sparse_mla_triton
from deepinfer.h20_ops.sparse_mla_sm90 import sparse_mla_fwd_attn as tl_sparse_mla

DEV = "cuda"
HQ = 64
D = 512
FAILED: list[str] = []


def check(tag: str, cond: bool, detail: str = "") -> None:
    print(f"  {'ok  ' if cond else 'FAIL'} {tag}{'  ' + detail if detail else ''}")
    if not cond:
        FAILED.append(tag)


def close(tag: str, a: torch.Tensor, b: torch.Tensor) -> None:
    af, bf = a.float(), b.float()
    if af.shape != bf.shape:
        check(tag, False, f"shape {tuple(af.shape)} vs {tuple(bf.shape)}")
        return
    diff = (af - bf).abs().max().item()
    scale = af.abs().max().item()
    tol = 2.0**-8 * max(scale, 1e-9) * 2
    check(tag, diff <= tol, f"max diff {diff:.3e} tol {tol:.3e}")


def test_prefill_contract() -> None:
    """Shared-KV convention: the shape prefill actually uses."""
    print("\n[prefill: shared KV, batch=1]")
    torch.manual_seed(0)
    sq, skv, topk = 64, 2048, 512
    q = (torch.randn((1, sq, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    kv = (torch.randn((1, skv, 1, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    idx = torch.stack([torch.randperm(skv, device=DEV)[:topk] for _ in range(sq)])
    idx = idx.view(1, sq, 1, topk).to(torch.int32).contiguous()
    scale = D**-0.5

    tl_out, tl_lse = tl_sparse_mla(q, kv, idx, sm_scale=scale, is_causal=False)
    tri_out, tri_lse = sparse_mla_triton.sparse_mla_fwd_attn(q, kv, idx, sm_scale=scale, is_causal=False)
    torch.cuda.synchronize()

    check("returns exactly two tensors", isinstance((tri_out, tri_lse), tuple) and len((tri_out, tri_lse)) == 2)
    check("out dtype matches reference", tri_out.dtype == tl_out.dtype, str(tri_out.dtype))
    check("lse dtype matches reference", tri_lse.dtype == tl_lse.dtype, str(tri_lse.dtype))
    check("out shape preserved", tri_out.shape == q.shape, str(tuple(tri_out.shape)))
    check("lse shape", tri_lse.shape == (1, sq, HQ), str(tuple(tri_lse.shape)))
    close("out", tl_out, tri_out)
    close("lse (log2)", tl_lse, tri_lse)


def test_lse_units() -> None:
    """lse must be log2, not the kernel's natural log.

    Asserted independently of the out comparison: an adapter that forgot the conversion would still
    produce the right `out` here (the kernel applies no sink), so comparing out alone cannot catch
    it. The ratio between the two units is exactly 1/ln2, so it is checkable without a reference.
    """
    print("\n[lse units: log2, not natural log]")
    torch.manual_seed(1)
    sq, skv, topk = 16, 1024, 512
    q = (torch.randn((1, sq, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    kv = (torch.randn((1, skv, 1, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    idx = torch.randint(0, skv, (1, sq, 1, topk), device=DEV, dtype=torch.int32).contiguous()

    _tl_out, tl_lse = tl_sparse_mla(q, kv, idx, sm_scale=D**-0.5, is_causal=False)
    _tri_out, tri_lse = sparse_mla_triton.sparse_mla_fwd_attn(q, kv, idx, sm_scale=D**-0.5, is_causal=False)
    torch.cuda.synchronize()

    # If the adapter had passed the natural log through, tri_lse would equal tl_lse * ln2 instead
    # of tl_lse, i.e. be smaller by a factor of 1.4427. Check the ratio is ~1, not ~0.693.
    ratio = (tri_lse.float() / tl_lse.float().clamp_min(1e-6)).median().item()
    check("lse ratio to reference is ~1.0 (would be 0.693 if unconverted)", abs(ratio - 1.0) < 0.02, f"ratio {ratio:.4f}")
    # And the natural-log interpretation must be measurably different, or the test is vacuous.
    natural_ratio = (tl_lse.float() * math.log(2) / tl_lse.float().clamp_min(1e-6)).median().item()
    check("the two units are distinguishable for this input", abs(natural_ratio - 1.0) > 0.2, f"ln/log2 = {natural_ratio:.3f}")


def test_decode_convention() -> None:
    """Per-query KV convention: what decode passes (batch axis is the query axis)."""
    print("\n[decode: per-query KV, batch=total_q]")
    torch.manual_seed(2)
    total_q, window = 32, 512
    q = (torch.randn((total_q, 1, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    # Every query gets a different window, so a mis-mapped convention shows up as wrong numbers.
    kv = (torch.randn((total_q, window, 1, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    identity = torch.arange(window, dtype=torch.int32, device=DEV).view(1, 1, 1, -1)
    idx = identity.expand(total_q, 1, 1, window).contiguous()

    tl_out, tl_lse = tl_sparse_mla(q, kv, idx, sm_scale=D**-0.5, is_causal=False)
    tri_out, tri_lse = sparse_mla_triton.sparse_mla_fwd_attn(q, kv, idx, sm_scale=D**-0.5, is_causal=False)
    torch.cuda.synchronize()

    close("decode out", tl_out.view(total_q, HQ, D), tri_out.view(total_q, HQ, D))
    close("decode lse", tl_lse.view(total_q, HQ), tri_lse.view(total_q, HQ))


def test_rejects_unsupported() -> None:
    """Loud failure beats a silently wrong result."""
    print("\n[rejects what it cannot do]")
    q = torch.zeros((1, 4, HQ, D), device=DEV, dtype=torch.bfloat16)
    kv = torch.zeros((1, 8, 1, D), device=DEV, dtype=torch.bfloat16)
    idx = torch.zeros((1, 4, 1, 8), device=DEV, dtype=torch.int32)

    try:
        sparse_mla_triton.sparse_mla_fwd_attn(q, kv, idx, is_causal=True)
        check("is_causal=True raises", False)
    except NotImplementedError as exc:
        check("is_causal=True raises", "is_causal=False" in str(exc))

    q2 = torch.zeros((1, 4, 32, D), device=DEV, dtype=torch.bfloat16)
    try:
        sparse_mla_triton.sparse_mla_fwd_attn(q2, kv, idx, is_causal=False)
        check("unsupported head count raises", False)
    except ValueError as exc:
        check("unsupported head count raises", "heads" in str(exc))

    kv3 = torch.zeros((3, 8, 1, D), device=DEV, dtype=torch.bfloat16)
    try:
        sparse_mla_triton.sparse_mla_fwd_attn(q, kv3, idx, is_causal=False)
        check("mismatched kv batch raises", False)
    except ValueError as exc:
        check("mismatched kv batch raises", "batch axis" in str(exc))


def test_match_sink_behaviour() -> None:
    """The consumer of lse is `_apply_sink`; check it two ways rather than one.

    Comparing only the final bf16 tensors is the wrong instrument. `_apply_sink` multiplies an
    already-bf16 `out` (itself differing by ~1.4 ULP) by a factor and rounds again, so the result can
    differ by a couple ULP even when both inputs are correct -- measured 1.83e-4 against 1.76e-4 for
    2 ULP of the input magnitude. That is a tolerance artifact, not a disagreement, so this asserts
    the thing that actually matters on its own: the multiplicative factor, in fp32, is a pure
    function of lse and must agree tightly. That is precisely the claim "lse units are right".
    """
    print("\n[attn_sink path agrees]")
    from deepinfer.h20_ops.attn_sm90 import _apply_sink

    torch.manual_seed(3)
    sq, skv, topk = 16, 1024, 512
    q = (torch.randn((1, sq, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    kv = (torch.randn((1, skv, 1, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    idx = torch.randint(0, skv, (1, sq, 1, topk), device=DEV, dtype=torch.int32).contiguous()
    sink = torch.randn((HQ,), device=DEV, dtype=torch.float32) * 0.5

    tl_out, tl_lse = tl_sparse_mla(q, kv, idx, sm_scale=D**-0.5, is_causal=False)
    tri_out, tri_lse = sparse_mla_triton.sparse_mla_fwd_attn(q, kv, idx, sm_scale=D**-0.5, is_causal=False)
    torch.cuda.synchronize()

    f_tl = torch.sigmoid(tl_lse.float() * math.log(2) - sink.view(1, 1, HQ)).unsqueeze(-1)
    f_tri = torch.sigmoid(tri_lse.float() * math.log(2) - sink.view(1, 1, HQ)).unsqueeze(-1)
    rel = (f_tl - f_tri).abs().max().item() / max(f_tl.abs().max().item(), 1e-9)
    check("sink factor relative diff < 1e-5", rel < 1e-5, f"rel {rel:.3e}")

    # Composed path: two bf16 roundings, so 4 ULP of the input magnitude, not 2.
    a = _apply_sink(tl_out, tl_lse, sink).float()
    b = _apply_sink(tri_out, tri_lse, sink).float()
    diff = (a - b).abs().max().item()
    scale = tl_out.float().abs().max().item()
    tol = 4 * 2.0**-8 * max(scale, 1e-9)
    check("_apply_sink output within 4 ULP", diff <= tol, f"max diff {diff:.3e} tol {tol:.3e}")


TESTS = [
    test_prefill_contract,
    test_lse_units,
    test_decode_convention,
    test_rejects_unsupported,
    test_match_sink_behaviour,
]

if __name__ == "__main__":
    for test in TESTS:
        test()
    print()
    if FAILED:
        print(f"FAILED {len(FAILED)}: {FAILED}")
        sys.exit(1)
    print(f"ALL_PASSED ({len(TESTS)} tests)")