"""
D1.3 SMEM-P: Debug why hd>64 fails.
Test: compute raw PV (before O normalization) at hd=128 with SMEM-P
and compare against FP32 oracle.
Also test: hd=64 with SMEM-P but skip O normalization to isolate the error.
"""
import torch, math
import cutlass, cutlass.cute as cute
from cutlass import Float32, BFloat16
import cutlass.torch as ct
import cuda.bindings.driver as cuda
from dsv4.kernels.attention.fmha import FmhaKernel


def test_unnormalized(hd, use_smem_p, s_k=128):
    """Test with normalize=False to get raw O + LSE, isolate the P write error."""
    pv_n = min(hd, 256)
    q = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda')
    k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda')
    v = torch.randn(s_k, pv_n, dtype=torch.bfloat16, device='cuda')
    c = torch.zeros(128, pv_n, 1, dtype=torch.bfloat16, device='cuda')
    lse = torch.zeros(1, dtype=torch.float32, device='cuda')

    qf = q[:, :, 0].float()
    kf = k[:, :, 0].float()
    vf = v.float()
    scale = 1.0 / math.sqrt(hd)
    attn = qf @ kf.T * scale
    attn_softmax = torch.softmax(attn, dim=-1)
    ref = attn_softmax @ vf  # normalized reference
    ref_unnorm = attn_softmax * attn_softmax.shape[-1]  # just for debugging

    kern = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, normalize=False)
    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

    v_tile = v.unsqueeze(-1)
    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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile))
    mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c))
    mLSE = ct.from_dlpack(lse).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse))

    mode = "SMEM-P" if use_smem_p else "TMEM-P"
    print(f'Compiling hd={hd} {mode} normalize=False...', flush=True)
    compiled = cute.compile(kern, mQ, mK, mV, mC, stream, mLSE)
    compiled(mQ, mK, mV, mC, stream, mLSE)
    torch.cuda.synchronize()

    out = c[:, :, 0].float()
    lse_val = lse.item()

    # The un-normalized output should be: O_unnorm = exp(lse) * O_norm
    # So O_norm = O_unnorm / exp(lse)
    if lse_val != 0 and not math.isnan(lse_val) and not math.isinf(lse_val):
        out_norm = out / math.exp(lse_val)
        cos = torch.nn.functional.cosine_similarity(
            out_norm.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)
        ).item()
        max_abs = (out_norm - ref).abs().max().item()
        print(f'  hd={hd} {mode} unnorm: cos={cos:.6f}  max_abs={max_abs:.6f}')
        print(f'    LSE={lse_val:.6f}  exp(lse)={math.exp(lse_val):.6f}')
        print(f'    out range: [{out.min().item():.4f}, {out.max().item():.4f}]')
        print(f'    ref range: [{ref.min().item():.4f}, {ref.max().item():.4f}]')
    else:
        print(f'  hd={hd} {mode} unnorm: INVALID LSE={lse_val}')
        print(f'    out has NaN: {torch.isnan(out).any().item()}')
        print(f'    out range: [{out.min().item():.4f}, {out.max().item():.4f}]')

    # Also test normalized
    kern2 = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, normalize=True)
    c2 = torch.zeros(128, pv_n, 1, dtype=torch.bfloat16, device='cuda')
    mC2 = ct.from_dlpack(c2).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c2))
    compiled2 = cute.compile(kern2, mQ, mK, mV, mC2, stream)
    compiled2(mQ, mK, mV, mC2, stream)
    torch.cuda.synchronize()

    out2 = c2[:, :, 0].float()
    cos2 = torch.nn.functional.cosine_similarity(
        out2.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)
    ).item()
    print(f'  hd={hd} {mode} normalized: cos={cos2:.6f}')
    print()


if __name__ == '__main__':
    print("=== SMEM-P Debug: Unnormalized vs Normalized ===\n")

    # hd=64 baseline
    test_unnormalized(64, use_smem_p=False)
    test_unnormalized(64, use_smem_p=True)

    # hd=128
    test_unnormalized(128, use_smem_p=True)

    # hd=256
    test_unnormalized(256, use_smem_p=True)