"""
FMHA D3: In-kernel SWA sequence length masking.

Proper approach: the kernel receives swa_lens and masks logits to -inf
inside the softmax, using the tTMEM_LOADcS coordinate tensor to map
register fragment positions to (row, col) in the QK matrix.

This replaces the pre-masking approach (BF16 min on K) which cannot
produce true -inf QK scores.

Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d3_inkernel_mask.py
"""
import torch
import math
import cutlass.cute as cute
import cutlass.torch as ct
import cuda.bindings.driver as cuda
from dsv4.kernels.attention.fmha import FmhaKernel


def reference_swa_attention(q, k, v, swa_lens, scale):
    """FP32 reference with proper -inf masking.

    Args:
        q: (M, hd) BF16
        k: (s_k, hd) BF16
        v: (s_k, hd) BF16
        swa_lens: (M,) int32 — per-row number of valid KV positions
        scale: float

    Returns:
        o: (M, hd) BF16
    """
    scores = torch.matmul(q.float(), k.float().T) * scale
    for i in range(q.shape[0]):
        sl = swa_lens[i].item()
        if sl < k.shape[0]:
            scores[i, sl:] = float('-inf')
    max_s = scores.max(dim=-1, keepdim=True).values
    exp_s = (scores - max_s).exp()
    sum_s = exp_s.sum(dim=-1, keepdim=True)
    p = exp_s / sum_s
    o = torch.matmul(p, v.float())
    return o.to(torch.bfloat16)


def _run_fmha_masked(q_3d, k_3d, v, m, s_k, hd, swa_lens_tensor, use_smem_p=False):
    """Run FMHA with in-kernel SWA masking and return normalized output.

    Args:
        q_3d: (M, hd, 1) BF16
        k_3d: (s_k, hd, 1) BF16
        v: (s_k, hd) BF16
        swa_lens_tensor: (1,) int32 — number of valid KV positions

    Returns:
        o_norm: (M, hd) BF16
    """
    scale = 1.0 / math.sqrt(hd)
    kernel = FmhaKernel(
        head_dim=hd, s_k=s_k, use_smem_p=use_smem_p,
        apply_swa_mask=True, normalize=False,
    )
    pv_n_tile = kernel.pv_n_tile
    n_pv_tiles = kernel.n_pv_tiles
    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

    # swa_lens as CuTe tensor (1D, int32)
    mSwaLens = ct.from_dlpack(swa_lens_tensor).mark_layout_dynamic(
        leading_dim=ct.get_leading_dim(swa_lens_tensor)
    )

    o_unnorm = torch.zeros(m, hd, dtype=torch.float32, device='cuda')
    lse_all = torch.zeros(m, dtype=torch.float32, device='cuda')

    for pv in range(n_pv_tiles):
        v_tile = v[:, pv * pv_n_tile:(pv + 1) * pv_n_tile].contiguous().unsqueeze(-1)
        c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
        lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')

        mQ = ct.from_dlpack(q_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_3d))
        mK = ct.from_dlpack(k_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_3d))
        mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile))
        mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile))
        mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))

        if pv == 0:
            compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE, mSwaLens)

        compiled(mQ, mK, mV, mC, stream, mLSE, mSwaLens)
        o_unnorm[:, pv * pv_n_tile:(pv + 1) * pv_n_tile] = c_tile[:, :, 0].float()
        lse_all += lse_tensor[:, 0, 0]

    # External normalization using LSE
    # O_norm = O_unnorm / exp(LSE) ... but per-row LSE only row 0 is written.
    # Use reference attn_sum for normalization (same as head-packed tests).
    q_flat = q_3d[:, :, 0]
    k_flat = k_3d[:, :, 0]
    scores = torch.matmul(q_flat.float(), k_flat.float().T) * scale
    swa_len_val = swa_lens_tensor[0].item()
    if swa_len_val < s_k:
        scores[:, swa_len_val:] = float('-inf')
    max_s = scores.max(dim=-1, keepdim=True).values
    attn_sum = (scores - max_s).exp().sum(dim=-1, keepdim=True)
    o_norm = (o_unnorm / attn_sum).to(torch.bfloat16)
    return o_norm


def test_d3_no_mask():
    """Full window (swa_lens=128): no masking, regression test."""
    print("\n=== Test 1: No masking (swa_lens=128, hd=64) ===")
    torch.manual_seed(42)
    m, s_k, hd = 128, 128, 64
    scale = 1.0 / math.sqrt(hd)

    q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
    k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda')
    v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
    swa_lens = torch.tensor([s_k], dtype=torch.int32, device='cuda')

    o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_lens)
    ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, swa_lens.cpu().expand(m), scale)

    cos = torch.nn.functional.cosine_similarity(
        o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0)
    ).item()
    print(f"  cos = {cos:.6f}")
    assert cos >= 0.995, f"Regression: cosine too low: {cos}"
    print("  ✅ PASS")


def test_d3_swa64():
    """SWA with swa_lens=64: mask positions 64-127 to -inf."""
    print("\n=== Test 2: swa_lens=64 (hd=64) ===")
    torch.manual_seed(42)
    m, s_k, hd = 128, 128, 64
    scale = 1.0 / math.sqrt(hd)

    q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
    k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda')
    v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
    swa_lens = torch.tensor([64], dtype=torch.int32, device='cuda')

    o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_lens)
    ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, swa_lens.cpu().expand(m), scale)

    cos = torch.nn.functional.cosine_similarity(
        o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0)
    ).item()
    print(f"  cos = {cos:.6f}")
    assert cos >= 0.99, f"cosine too low: {cos}"
    print("  ✅ PASS")


def test_d3_swa32():
    """SWA with swa_lens=32: only 32 valid positions."""
    print("\n=== Test 3: swa_lens=32 (hd=64) ===")
    torch.manual_seed(42)
    m, s_k, hd = 128, 128, 64
    scale = 1.0 / math.sqrt(hd)

    q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
    k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda')
    v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
    swa_lens = torch.tensor([32], dtype=torch.int32, device='cuda')

    o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_lens)
    ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, swa_lens.cpu().expand(m), scale)

    cos = torch.nn.functional.cosine_similarity(
        o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0)
    ).item()
    print(f"  cos = {cos:.6f}")
    assert cos >= 0.99, f"cosine too low: {cos}"
    print("  ✅ PASS")


def test_d3_swa1():
    """Edge case: swa_lens=1, only one valid KV position."""
    print("\n=== Test 4: swa_lens=1 (hd=64) ===")
    torch.manual_seed(42)
    m, s_k, hd = 128, 128, 64
    scale = 1.0 / math.sqrt(hd)

    q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
    k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda')
    v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
    swa_lens = torch.tensor([1], dtype=torch.int32, device='cuda')

    o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_lens)
    ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, swa_lens.cpu().expand(m), scale)

    cos = torch.nn.functional.cosine_similarity(
        o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0)
    ).item()
    print(f"  cos = {cos:.6f}")
    assert cos >= 0.99, f"cosine too low: {cos}"
    print("  ✅ PASS")


def test_d3_hd128():
    """SWA masking at hd=128 (SMEM-P path)."""
    print("\n=== Test 5: swa_lens=64 (hd=128) ===")
    torch.manual_seed(42)
    m, s_k, hd = 128, 128, 128
    scale = 1.0 / math.sqrt(hd)

    q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
    k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda')
    v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
    swa_lens = torch.tensor([64], dtype=torch.int32, device='cuda')

    o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_lens, use_smem_p=True)
    ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, swa_lens.cpu().expand(m), scale)

    cos = torch.nn.functional.cosine_similarity(
        o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0)
    ).item()
    print(f"  cos = {cos:.6f}")
    assert cos >= 0.99, f"cosine too low: {cos}"
    print("  ✅ PASS")


def test_d3_swa128_hd128():
    """No masking at hd=128: regression test (should match existing D1 results)."""
    print("\n=== Test 6: No masking (swa_lens=128, hd=128) ===")
    torch.manual_seed(42)
    m, s_k, hd = 128, 128, 128
    scale = 1.0 / math.sqrt(hd)

    q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
    k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda')
    v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
    swa_lens = torch.tensor([s_k], dtype=torch.int32, device='cuda')

    o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_lens, use_smem_p=True)
    ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, swa_lens.cpu().expand(m), scale)

    cos = torch.nn.functional.cosine_similarity(
        o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0)
    ).item()
    print(f"  cos = {cos:.6f}")
    assert cos >= 0.995, f"Regression: cosine too low: {cos}"
    print("  ✅ PASS")


def test():
    print("=== D3: In-Kernel SWA Sequence Length Mask ===")
    test_d3_no_mask()
    test_d3_swa64()
    test_d3_swa32()
    test_d3_swa1()
    test_d3_hd128()
    test_d3_swa128_hd128()
    print("\n=== ALL TESTS PASSED ===")


if __name__ == '__main__':
    test()