"""
FMHA D3: SWA sequence length mask (large-negative pre-masking approach).

K/V rows at positions >= swa_lens are set to BF16 min (-65504) before
passing to the kernel. This gives very large negative QK scores for
invalid positions, producing exp(score) ≈ 0 contribution to the softmax.
Effectively equivalent to -inf masking for practical purposes.

Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d3_swa_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

BF16_MIN = torch.tensor(-65504.0, dtype=torch.bfloat16)


def reference_swa_attention(q, k, v, swa_lens, scale):
    """FP32 reference with proper -inf masking."""
    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), sum_s


def _run_fmha(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False):
    """Run FMHA 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)
    pv_n_tile = kernel.pv_n_tile; n_pv_tiles = kernel.n_pv_tiles
    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
    
    v_tile = v[:, 0: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))
    
    compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE)
    
    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.zero_(); lse_tensor.zero_()
        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))
        compiled(mQ, mK, mV, mC, stream, mLSE)
        o_unnorm[:, pv*pv_n_tile:(pv+1)*pv_n_tile] = c_tile[:,:,0].float()
    
    # Reference normalization
    q_flat = q_3d[:,:,0]; k_flat = k_3d[:,:,0]
    scores = torch.matmul(q_flat.float(), k_flat.float().T) * scale
    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_full_window():
    """Full SWA window (swa_lens=128): no masking needed."""
    print("\n=== Test 1: Full SWA window (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')
    
    o = _run_fmha(q, k, v, m, s_k, hd)
    ref, _ = reference_swa_attention(q[:,:,0], k[:,:,0], v, torch.full((m,), s_k, dtype=torch.int32, device='cuda'), 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
    print("  ✅ PASS")


def test_d3_swa64():
    """SWA with swa_lens=64: mask K rows >= 64 with BF16 min."""
    print("\n=== Test 2: SWA 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.full((m,), 64, dtype=torch.int32, device='cuda')
    
    # Mask K rows >= 64 with BF16 min
    k_masked = k.clone()
    k_masked[64:] = BF16_MIN.to(k.device)
    # Also mask V (otherwise invalid positions contribute to output)
    v_masked = v.clone()
    v_masked[64:] = 0
    
    o = _run_fmha(q, k_masked, v_masked, m, s_k, hd)
    ref, _ = reference_swa_attention(q[:,:,0], k[:,:,0], v, swa_lens, 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 tokens."""
    print("\n=== Test 3: SWA 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.full((m,), 32, dtype=torch.int32, device='cuda')
    
    k_masked = k.clone()
    k_masked[32:] = BF16_MIN.to(k.device)
    v_masked = v.clone()
    v_masked[32:] = 0
    
    o = _run_fmha(q, k_masked, v_masked, m, s_k, hd)
    ref, _ = reference_swa_attention(q[:,:,0], k[:,:,0], v, swa_lens, 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():
    print("=== D3: SWA Sequence Length Mask ===")
    test_d3_full_window()
    test_d3_swa64()
    test_d3_swa32()
    print("\n=== ALL TESTS PASSED ===")


if __name__ == '__main__':
    test()