"""
D1: Test TMEM round-trip on O in isolation.

Runs the kernel with s_k=128 (1 KV tile, no rescale needed).
Then manually does a load-modify-store round-trip on O in TMEM
using the correction_rescale atoms.
If the round-trip corrupts data, we know the atoms are broken.
If it preserves data, the bug is elsewhere.
"""
import torch, 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 test():
    hd = 64
    s_k = 128  # 1 KV tile, no rescale needed
    m = 128
    torch.manual_seed(42)

    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')
    c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')

    # FP32 reference
    qf = q[:, :, 0].float()
    kf = k[:, :, 0].float()
    scale = 1.0 / math.sqrt(hd)
    attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0]
    attn_exp = torch.exp(qf @ kf.T * scale - attn_max)
    attn_sum = attn_exp.sum(dim=-1, keepdim=True)
    ref_unnorm = attn_exp @ v.float()

    lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')

    # Test 1: s_k=128 baseline (no rescale) — should be PASS
    kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False)
    pv_n_tile = kernel.pv_n_tile
    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

    v_tile = v[:, 0:pv_n_tile].contiguous()
    v_kernel = v_tile.unsqueeze(-1)
    c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')

    mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
    mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
    mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
    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))

    print(f'Test 1: s_k=128 baseline (no rescale)', flush=True)
    compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE)
    compiled(mQ, mK, mV, mC, stream, mLSE)
    torch.cuda.synchronize()

    out1 = c_tile[:, :, 0].float()
    cos1 = torch.nn.functional.cosine_similarity(
        out1.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0)
    ).item()
    print(f'  cos_unnorm={cos1:.6f}  {"PASS" if cos1 >= 0.99 else "FAIL"}')

    # Test 2: s_k=256 with rescale — this is the failing test
    s_k2 = 256
    k2 = torch.randn(s_k2, hd, 1, dtype=torch.bfloat16, device='cuda')
    v2 = torch.randn(s_k2, hd, dtype=torch.bfloat16, device='cuda')
    c2 = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')

    kf2 = k2[:, :, 0].float()
    attn_max2 = (qf @ kf2.T * scale).max(dim=-1, keepdim=True)[0]
    attn_exp2 = torch.exp(qf @ kf2.T * scale - attn_max2)
    attn_sum2 = attn_exp2.sum(dim=-1, keepdim=True)
    ref_unnorm2 = attn_exp2 @ v2.float()

    kernel2 = FmhaKernel(head_dim=hd, s_k=s_k2, use_smem_p=False, normalize=False)
    lse_tensor2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')

    v_tile2 = v2[:, 0:pv_n_tile].contiguous()
    v_kernel2 = v_tile2.unsqueeze(-1)
    c_tile2 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')

    mK2 = ct.from_dlpack(k2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k2))
    mV2 = ct.from_dlpack(v_kernel2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel2))
    mC2 = ct.from_dlpack(c_tile2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile2))
    mLSE2 = ct.from_dlpack(lse_tensor2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor2))

    print(f'Test 2: s_k=256 with O rescale', flush=True)
    compiled2 = cute.compile(kernel2, mQ, mK2, mV2, mC2, stream, mLSE2)
    compiled2(mQ, mK2, mV2, mC2, stream, mLSE2)
    torch.cuda.synchronize()

    out2 = c_tile2[:, :, 0].float()
    cos2 = torch.nn.functional.cosine_similarity(
        out2.flatten().unsqueeze(0), ref_unnorm2.flatten().unsqueeze(0)
    ).item()
    print(f'  cos_unnorm={cos2:.6f}  {"PASS" if cos2 >= 0.99 else "FAIL"}')


if __name__ == '__main__':
    test()