"""CPU correctness tests for the pure-torch flash_mla vendor package.

Run: python -I test_flash_mla.py   (PYTHONPATH must contain the vendor dir)
"""

import sys

import torch

import flash_mla
from flash_mla import flash_mla_sparse_fwd, flash_mla_with_kvcache, get_mla_metadata
from flash_mla.flash_mla_interface import FlashMLASchedMeta

import ref_attn
import ref_quant

torch.manual_seed(0)
FAILS = []


def check(name, got, exp, atol=2e-2, rtol=2e-2, exact=False):
    assert got.shape == exp.shape, f"{name}: shape {tuple(got.shape)} != {tuple(exp.shape)}"
    g, e = got.float(), exp.float()
    if exact:
        ok = torch.equal(g, e)
        msg = "bitwise"
    else:
        same_inf = ((g == e) | ((g != g) & (e != e))).float()
        finite = torch.isfinite(g) & torch.isfinite(e)
        ok = bool(torch.all(torch.where(finite, (g - e).abs() <= atol + rtol * e.abs(), same_inf.bool())))
        msg = f"max abs diff {(g - e).abs()[finite].max().item() if finite.any() else 0.0:.3e}"
    print(f"[{'ok ' if ok else 'FAIL'}] {name} ({msg})")
    if not ok:
        FAILS.append(name)


def make_indices(rows, topk, s_kv, invalid_frac=0.2, invalid_value=-1):
    idx = torch.randint(0, s_kv, (rows, topk), dtype=torch.int32)
    mask = torch.rand(rows, topk) < invalid_frac
    idx = torch.where(mask, torch.full_like(idx, invalid_value), idx)
    return idx


# ---------------------------------------------------------------- prefill
def test_prefill(h_q=64, s_q=7, s_kv=41, topk=24, with_sink=True, with_len=True, d=512):
    q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
    kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
    idx = make_indices(s_q, topk, s_kv).unsqueeze(1)
    # also exercise the ">= s_kv" invalid convention of the prefill kernel
    idx[0, 0, 0] = s_kv + 3
    sink = torch.randn(h_q) if with_sink else None
    tl = torch.randint(1, topk + 1, (s_q,), dtype=torch.int32) if with_len else None
    sm = d**-0.5

    out, ml, lse = flash_mla_sparse_fwd(q, kv, idx, sm, 512, sink, tl)
    r_out, r_ml, r_lse = ref_attn.ref_sparse_attn_fwd(q, kv, idx, sm, tl, sink)
    tag = f"prefill h{h_q} sink={with_sink} len={with_len}"
    check(f"{tag} out", out, r_out)
    check(f"{tag} max_logits", ml, r_ml)
    check(f"{tag} lse", lse, r_lse)
    assert out.dtype == torch.bfloat16 and lse.dtype == torch.float32 and ml.dtype == torch.float32


def test_prefill_lonely():
    h_q, s_q, s_kv, topk, d = 64, 3, 16, 8, 512
    q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
    kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
    idx = torch.full((s_q, 1, topk), -1, dtype=torch.int32)
    out, ml, lse = flash_mla_sparse_fwd(q, kv, idx, d**-0.5)
    check("prefill lonely out==0", out, torch.zeros_like(out), exact=True)
    check("prefill lonely lse==+inf", lse, torch.full_like(lse, float("inf")), exact=True)
    check("prefill lonely max_logits==-inf", ml, torch.full_like(ml, float("-inf")), exact=True)


# ---------------------------------------------------------------- kv cache
def test_kv_dequant():
    for fmt in ("v41_fp8", "v41_fp4"):
        raw = (torch.randn(5, 8, 1, 512) / 4).to(torch.bfloat16)
        cache = ref_quant.quantize_k_cache(raw, fmt)
        exp = ref_quant.dequantize_k_cache(cache, fmt).view(5, 8, 512)
        got = flash_mla.kv_cache.dequantize_records(cache.view(5, 8, -1), fmt)
        check(f"dequant {fmt}", got, exp, exact=True)
        # relative error against the bf16 source
        err = ((got.float() - raw.view(5, 8, 512).float()).abs() / raw.view(5, 8, 512).float().abs().clamp_min(1e-3)).median()
        print(f"       {fmt} median rel err vs bf16 source: {err.item():.4f}")
        # byte-view equivalence: int8 / float8 views of the same bytes must dequantize identically
        got_i8 = flash_mla.kv_cache.dequantize_records(cache.view(5, 8, -1).view(torch.int8), fmt)
        check(f"dequant {fmt} int8 view", got_i8, exp, exact=True)


# ---------------------------------------------------------------- decode
def test_decode(h_q=64, b=5, topk=20, block=4, nblocks=9, with_sink=True, with_len=True, extra=False, strided=False):
    d, s_q = 512, 1
    raw = (torch.randn(nblocks, block, 1, d) / 8).to(torch.bfloat16)
    cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
    if strided:
        # non-contiguous block stride, like deepinfer/worker/cache_engine's as_strided views
        pad = torch.zeros(nblocks, block * 528 + 64, dtype=torch.uint8)
        pad[:, : block * 528] = cache.reshape(nblocks, -1)
        cache = pad.as_strided((nblocks, block, 1, 528), (pad.stride(0), 528, 528, 1))
    kv_bf16 = ref_quant.dequantize_k_cache(cache, "v41_fp8").view(-1, d)

    q = (torch.randn(b, s_q, h_q, d) / 8).to(torch.bfloat16)
    idx = make_indices(b * s_q, topk, nblocks * block).view(b, s_q, topk)
    sink = torch.randn(h_q) if with_sink else None
    tl = torch.randint(1, topk + 1, (b,), dtype=torch.int32) if with_len else None
    sm = d**-0.5

    scopes = [(kv_bf16, idx, tl)]
    e_cache = e_idx = e_tl = None
    if extra:
        e_raw = (torch.randn(6, block, 1, d) / 8).to(torch.bfloat16)
        e_cache = ref_quant.quantize_k_cache(e_raw, "v41_fp4")
        e_kv = ref_quant.dequantize_k_cache(e_cache, "v41_fp4").view(-1, d)
        e_topk = 12
        e_idx = make_indices(b * s_q, e_topk, 6 * block).view(b, s_q, e_topk)
        e_tl = torch.randint(1, e_topk + 1, (b,), dtype=torch.int32)
        scopes.append((e_kv, e_idx, e_tl))

    meta, none = get_mla_metadata()
    assert none is None and isinstance(meta, FlashMLASchedMeta)
    out, lse = flash_mla_with_kvcache(
        q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
        tile_scheduler_metadata=meta, softmax_scale=sm, causal=False, is_fp8_kvcache=True,
        indices=idx, attn_sink=sink, extra_k_cache=e_cache, extra_indices_in_kvcache=e_idx,
        topk_length=tl, extra_topk_length=e_tl,
    )
    r_out, r_lse = ref_attn.ref_sparse_attn_decode(q, scopes, sm, sink)
    tag = f"decode h{h_q} sink={with_sink} len={with_len} extra={extra} strided={strided}"
    check(f"{tag} out", out, r_out)
    check(f"{tag} lse", lse, r_lse)
    assert meta.have_initialized and meta.config.topk == topk and meta.config.page_block_size == block
    # reuse with a different topk must be rejected
    try:
        flash_mla_with_kvcache(
            q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
            tile_scheduler_metadata=meta, softmax_scale=sm, indices=idx[..., :-1],
        )
        print("[FAIL] sched_meta reuse check")
        FAILS.append("sched_meta reuse")
    except AssertionError:
        print("[ok ] sched_meta reuse rejected on topk change")


def test_decode_multi_token():
    """s_q > 1 (MTP/DFlash style): topk_length is per request, broadcast over s_q."""
    d, h_q, b, s_q, topk, block, nblocks = 512, 64, 3, 4, 16, 8, 7
    raw = (torch.randn(nblocks, block, 1, d) / 8).to(torch.bfloat16)
    cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
    kv_bf16 = ref_quant.dequantize_k_cache(cache, "v41_fp8").view(-1, d)
    q = (torch.randn(b, s_q, h_q, d) / 8).to(torch.bfloat16)
    idx = make_indices(b * s_q, topk, nblocks * block).view(b, s_q, topk)
    tl = torch.randint(1, topk + 1, (b,), dtype=torch.int32)
    sm = d**-0.5
    out, lse = flash_mla_with_kvcache(
        q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
        tile_scheduler_metadata=FlashMLASchedMeta(), softmax_scale=sm, indices=idx, topk_length=tl,
    )
    r_out, r_lse = ref_attn.ref_sparse_attn_decode(q, [(kv_bf16, idx, tl)], sm)
    check("decode s_q=4 out", out, r_out)
    check("decode s_q=4 lse", lse, r_lse)


def test_decode_lonely():
    d, h_q, b, topk, block, nblocks = 512, 64, 2, 8, 4, 4
    raw = (torch.randn(nblocks, block, 1, d) / 8).to(torch.bfloat16)
    cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
    q = (torch.randn(b, 1, h_q, d) / 8).to(torch.bfloat16)
    idx = torch.full((b, 1, topk), -1, dtype=torch.int32)
    out, lse = flash_mla_with_kvcache(
        q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
        tile_scheduler_metadata=FlashMLASchedMeta(), softmax_scale=d**-0.5, indices=idx,
        attn_sink=torch.randn(h_q),
    )
    check("decode lonely out==0", out, torch.zeros_like(out), exact=True)
    check("decode lonely lse==+inf", lse, torch.full_like(lse, float("inf")), exact=True)


def test_sink_infinities():
    d, h_q, s_q, s_kv, topk = 512, 64, 4, 20, 8
    q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
    kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
    idx = torch.randint(0, s_kv, (s_q, 1, topk), dtype=torch.int32)
    sink = torch.full((h_q,), float("-inf"))
    sink[: h_q // 2] = float("inf")
    out, _, lse = flash_mla_sparse_fwd(q, kv, idx, d**-0.5, 512, sink)
    base, _, base_lse = flash_mla_sparse_fwd(q, kv, idx, d**-0.5, 512, None)
    check("sink=+inf -> out 0", out[:, : h_q // 2], torch.zeros_like(out[:, : h_q // 2]), exact=True)
    check("sink=-inf -> unchanged", out[:, h_q // 2 :], base[:, h_q // 2 :], atol=1e-2, rtol=1e-2)
    check("sink does not change lse", lse, base_lse, exact=True)


def test_dense_unsupported():
    try:
        flash_mla.flash_attn_varlen_func(None, None, None, None, None, 0, 0)
        FAILS.append("dense should raise")
        print("[FAIL] dense entry point should raise")
    except RuntimeError as e:
        print(f"[ok ] dense entry point raises: {str(e)[:60]}...")


def test_fused_gated():
    import flash_mla.fused_norm_rope_attn_rope_cast as fused

    for fn, args in ((fused.permute_q_b_proj, ((None, None), 64, 512)), (fused.permute_wv_proj, ((None, None), 1, 512))):
        try:
            fn(*args)
            FAILS.append("permute should raise")
        except RuntimeError as e:
            print(f"[ok ] {fn.__name__} raises: {e}")

    # the fused attention core itself (norm + rope + attn + rope), checked against the test ref
    d, h_q, s_q, s_kv, topk = 512, 64, 5, 32, 16
    q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
    kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
    idx = make_indices(s_q, topk, s_kv).unsqueeze(1)
    positions = torch.randint(0, 1024, (s_q,), dtype=torch.int32)
    cache = torch.randn(1024, 64)
    rms = fused._rms_norm_scale_factor(q, 1e-4)
    q_roped = fused._rope(q, False, cache, positions, 64).to(torch.bfloat16)
    out, ml, lse = flash_mla_sparse_fwd(q_roped, kv, idx, d**-0.5, 512, None, None, _rms_norm_scale_factor=rms)
    out = fused._rope(out, True, cache, positions, 64)
    r_out, r_ml, r_lse = ref_attn.ref_sparse_attn_fwd(q_roped, kv, idx, d**-0.5, None, None, rms=rms)
    # reference O-RoPE, computed independently
    exp = r_out.float().clone()
    cos = cache[positions.long(), :32].unsqueeze(1)
    sin = -cache[positions.long(), 32:].unsqueeze(1)
    rp = exp[..., -64:]
    x0, x1 = rp[..., 0::2].clone(), rp[..., 1::2].clone()
    rp[..., 0::2] = x0 * cos - x1 * sin
    rp[..., 1::2] = x0 * sin + x1 * cos
    check("fused norm+rope+attn+rope out", out, exp, atol=3e-2, rtol=3e-2)
    check("fused lse", lse, r_lse)


def test_extra_fp8_cache():
    """SWA main cache + a second fp8 (528 B) cache as `extra`, the non-fp4 DeepInfer variant."""
    d, h_q, b, topk, block = 512, 64, 3, 12, 4
    raw = (torch.randn(5, block, 1, d) / 8).to(torch.bfloat16)
    cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
    kv = ref_quant.dequantize_k_cache(cache, "v41_fp8").view(-1, d)
    e_raw = (torch.randn(7, block, 1, d) / 8).to(torch.bfloat16)
    e_cache = ref_quant.quantize_k_cache(e_raw, "v41_fp8")
    e_kv = ref_quant.dequantize_k_cache(e_cache, "v41_fp8").view(-1, d)
    q = (torch.randn(b, 1, h_q, d) / 8).to(torch.bfloat16)
    idx = make_indices(b, topk, 5 * block).view(b, 1, topk)
    e_idx = make_indices(b, 9, 7 * block).view(b, 1, 9)
    out, lse = flash_mla_with_kvcache(
        q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
        tile_scheduler_metadata=FlashMLASchedMeta(), softmax_scale=d**-0.5, indices=idx,
        extra_k_cache=e_cache, extra_indices_in_kvcache=e_idx,
    )
    r_out, r_lse = ref_attn.ref_sparse_attn_decode(q, [(kv, idx, None), (e_kv, e_idx, None)], d**-0.5)
    check("decode extra fp8 cache out", out, r_out)
    check("decode extra fp8 cache lse", lse, r_lse)


def test_attention_backend_integration():
    """Drive deepinfer/model_executor/layers/attention_backend.py (head padding path included)."""
    import importlib.util
    import os

    path = os.path.join(os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..")), "deepinfer", "model_executor", "layers", "attention_backend.py")
    spec = importlib.util.spec_from_file_location("ab", path)
    ab = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(ab)

    d, h_q, s_q, s_kv, topk = 512, 40, 6, 30, 16  # 40 heads -> padded to 64
    q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
    kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
    idx = make_indices(s_q, topk, s_kv, invalid_value=-1)
    sink = torch.randn(h_q)
    ilen = torch.randint(1, topk + 1, (s_q,), dtype=torch.int32)
    out, lse = ab.prefill_attn_varlen_sparse(q, kv, idx, d**-0.5, attn_sink=sink, index_length=ilen)
    pq = torch.cat([q, q.new_zeros(s_q, 24, d)], dim=1)
    ps = torch.cat([sink, sink.new_zeros(24)])
    r_out, _, r_lse = ref_attn.ref_sparse_attn_fwd(pq, kv, idx.unsqueeze(1), d**-0.5, ilen, ps)
    check("attention_backend prefill out (head pad 40->64)", out, r_out[:, :h_q])
    check("attention_backend prefill lse", lse, r_lse[:, :h_q])

    b, block, nblocks = 4, 8, 6
    raw = (torch.randn(nblocks, block, 1, d) / 8).to(torch.bfloat16)
    cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
    kvb = ref_quant.dequantize_k_cache(cache, "v41_fp8").view(-1, d)
    qd = (torch.randn(b, h_q, d) / 8).to(torch.bfloat16)
    didx = make_indices(b, topk, nblocks * block).view(b, topk)
    dlen = torch.randint(1, topk + 1, (b,), dtype=torch.int32)
    out, lse = ab.decode_mla_mqa_sparse(
        qd, cache, 512, didx, FlashMLASchedMeta(), d**-0.5, attn_sink=sink, index_length=dlen
    )
    pqd = torch.cat([qd, qd.new_zeros(b, 24, d)], dim=1).view(b, 1, 64, d)
    r_out, r_lse = ref_attn.ref_sparse_attn_decode(pqd, [(kvb, didx.view(b, 1, topk), dlen)], d**-0.5, ps)
    check("attention_backend decode out (head pad)", out, r_out.view(b, 64, d)[:, :h_q])
    check("attention_backend decode lse", lse, r_lse.view(b, 64)[:, :h_q])


if __name__ == "__main__":
    test_kv_dequant()
    for h in (64, 128):
        test_prefill(h_q=h)
    test_prefill(with_sink=False, with_len=False)
    test_prefill_lonely()
    test_decode()
    test_decode(with_sink=False, with_len=False)
    test_decode(extra=True)
    test_decode(strided=True)
    test_decode(h_q=128)
    test_decode_multi_token()
    test_decode_lonely()
    test_sink_infinities()
    test_extra_fp8_cache()
    test_attention_backend_integration()
    test_dense_unsupported()
    test_fused_gated()
    print("\nFAILURES:", FAILS if FAILS else "none")
    sys.exit(1 if FAILS else 0)