"""
D1.5 Minimal TMEM round-trip test within FMHA.

Tests: run FMHA with s_k=128, then add a NOOP correction_rescale
(multiply by 1.0) BETWEEN the softmax and the epilogue.
If the output is bitwise identical to without rescale, the round-trip works.
If it differs, the round-trip corrupts data.

This test directly compares: same kernel, same data, WITH and WITHOUT
the O rescale step (forced to NOOP).
"""
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; m = 128; s_k = 128; scale = 1.0 / math.sqrt(hd)
    torch.manual_seed(42)
    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

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

    # Run WITHOUT rescale (s_k=128, n_kv_tiles=1, rescale code is dead)
    kernel1 = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False)
    pv_n_tile = kernel1.pv_n_tile
    c1 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
    lse1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
    rs1 = torch.zeros(m, 1, 1, dtype=torch.float32, 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))
    v_t = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1)
    mV = ct.from_dlpack(v_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_t))
    mC1 = ct.from_dlpack(c1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c1))
    mLSE1 = ct.from_dlpack(lse1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse1))
    mRS1 = ct.from_dlpack(rs1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs1))

    print('Compiling s_k=128 (no rescale)...', flush=True)
    comp1 = cute.compile(kernel1, mQ, mK, mV, mC1, stream, mLSE1, row_sums=mRS1)
    comp1(mQ, mK, mV, mC1, stream, mLSE1, row_sums=mRS1)
    torch.cuda.synchronize()

    # Run WITH NOOP rescale (s_k=256, but with debug_noop_rescale=True,
    # using same K/V repeated to make the second segment identical)
    # This triggers the O rescale code path (n_kv_tiles=2) but with factor=1.0
    k2 = torch.cat([k, k], dim=0)
    v2 = torch.cat([v, v], dim=0)

    kernel2 = FmhaKernel(head_dim=hd, s_k=256, use_smem_p=False, normalize=False, debug_noop_rescale=True)
    c2 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
    lse2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
    rs2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')

    mK2 = ct.from_dlpack(k2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k2))
    v2_t = v2[:, 0:pv_n_tile].contiguous().unsqueeze(-1)
    mV2 = ct.from_dlpack(v2_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v2_t))
    mC2 = ct.from_dlpack(c2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c2))
    mLSE2 = ct.from_dlpack(lse2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse2))
    mRS2 = ct.from_dlpack(rs2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs2))

    print('Compiling s_k=256 (NOOP rescale)...', flush=True)
    comp2 = cute.compile(kernel2, mQ, mK2, mV2, mC2, stream, mLSE2, row_sums=mRS2)
    comp2(mQ, mK2, mV2, mC2, stream, mLSE2, row_sums=mRS2)
    torch.cuda.synchronize()

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

    # Normalize both outputs
    out1 = c1[:, :, 0].float() / rs1[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30)
    out2 = c2[:, :, 0].float() / rs2[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30)

    cos1 = torch.nn.functional.cosine_similarity(out1.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0)).item()
    cos2 = torch.nn.functional.cosine_similarity(out2.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0)).item()
    cos_12 = torch.nn.functional.cosine_similarity(out1.flatten().unsqueeze(0), out2.flatten().unsqueeze(0)).item()

    # Also check unnormalized
    unnorm1 = c1[:, :, 0].float()
    unnorm2 = c2[:, :, 0].float()
    # With identical segments, s_k=256 unnorm should be 2x s_k=128 unnorm
    # (row_sum also doubles, so normalized result is the same)
    ratio = unnorm2 / unnorm1.clamp(min=1e-10)
    print(f'Unnorm ratio (should be ~2.0): mean={ratio.mean().item():.4f} std={ratio.std().item():.4f}')

    print(f's_k=128 (no rescale): cos={cos1:.6f}')
    print(f's_k=256 (NOOP rescale): cos={cos2:.6f}')
    print(f's_k=128 vs s_k=256: cos={cos_12:.6f}')

    # Bitwise comparison on unnormalized O
    # s_k=256 with identical segments: O_unnorm = 2 * O_unnorm_128
    expected_unnorm2 = 2.0 * unnorm1
    bit_cos = torch.nn.functional.cosine_similarity(unnorm2.flatten().unsqueeze(0), expected_unnorm2.flatten().unsqueeze(0)).item()
    max_diff = (unnorm2 - expected_unnorm2).abs().max().item()
    print(f'Unnorm: expected 2x, got cos={bit_cos:.6f} max_diff={max_diff:.4f}')


if __name__ == '__main__':
    test()