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

Proper approach: the kernel receives swa_len (int) 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.

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_len, scale):
    """FP32 reference with proper -inf masking."""
    scores = torch.matmul(q.float(), k.float().T) * scale
    if swa_len < k.shape[0]:
        scores[:, swa_len:] = 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_len_val, use_smem_p=False):
    """Run FMHA with in-kernel SWA masking and return normalized output."""
    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)

    o_unnorm = torch.zeros(m, hd, 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, swa_len_val)

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

    # External normalization using reference attn_sum
    q_flat = q_3d[:, :, 0]
    k_flat = k_3d[:, :, 0]
    scores = torch.matmul(q_flat.float(), k_flat.float().T) * scale
    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_len=128): no masking, regression test."""
    print("\n=== Test 1: No masking (swa_len=128, hd=64) ===")
    torch.manual_seed(42)
    m, s_k, hd = 128, 128, 64

    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')

    o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=s_k)
    ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, s_k, 1.0 / math.sqrt(hd))

    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_len=64: mask positions 64-127 to -inf."""
    print("\n=== Test 2: swa_len=64 (hd=64) ===")
    torch.manual_seed(42)
    m, s_k, hd = 128, 128, 64

    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')

    o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=64)
    ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, 64, 1.0 / math.sqrt(hd))

    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_len=32: only 32 valid positions."""
    print("\n=== Test 3: swa_len=32 (hd=64) ===")
    torch.manual_seed(42)
    m, s_k, hd = 128, 128, 64

    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')

    o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=32)
    ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, 32, 1.0 / math.sqrt(hd))

    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_len=1, only one valid KV position."""
    print("\n=== Test 4: swa_len=1 (hd=64) ===")
    torch.manual_seed(42)
    m, s_k, hd = 128, 128, 64

    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')

    o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=1)
    ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, 1, 1.0 / math.sqrt(hd))

    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_len=64 (hd=128) ===")
    torch.manual_seed(42)
    m, s_k, hd = 128, 128, 128

    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')

    o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=64, use_smem_p=True)
    ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, 64, 1.0 / math.sqrt(hd))

    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."""
    print("\n=== Test 6: No masking (swa_len=128, hd=128) ===")
    torch.manual_seed(42)
    m, s_k, hd = 128, 128, 128

    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')

    o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=s_k, use_smem_p=True)
    ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, s_k, 1.0 / math.sqrt(hd))

    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()