"""
D1.5 Debug: NO-OP TMEM round-trip test.
Tests s_k=256 with rescale_factor=1.0 (NO-OP).
If the round-trip itself is broken, even NO-OP will corrupt O.
If it produces the same (wrong) result as without rescale, the round-trip works.
"""
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 reference_attention(q, k, v, scale):
    qf = q.float(); kf = k.float()
    attn = qf @ kf.T * scale
    attn_max = attn.max(dim=-1, keepdim=True)[0]
    attn_exp = torch.exp(attn - attn_max)
    return attn_exp @ v.float()


def run_test(s_k, debug_noop=False, label=""):
    hd = 64; m = 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')

    kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False, debug_noop_rescale=debug_noop)
    pv_n_tile = kernel.pv_n_tile
    c = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
    lse = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
    rs = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')

    v_t = v[:, 0:pv_n_tile].contiguous().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_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_t))
    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))
    mRS = ct.from_dlpack(rs).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs))

    compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE, row_sums=mRS)
    compiled(mQ, mK, mV, mC, stream, mLSE, row_sums=mRS)
    torch.cuda.synchronize()

    ref = reference_attention(q[:, :, 0], k[:, :, 0], v, scale)
    out = c[:, :, 0].float()
    cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item()
    print(f'{label}: cos={cos:.6f} {"PASS" if cos >= 0.999 else "FAIL"}', flush=True)
    return out


def test():
    # Test 1: s_k=128 baseline
    run_test(128, label="s_k=128 baseline")

    # Test 2: s_k=256 WITH rescale (should be correct if TMEM round-trip works)
    run_test(256, label="s_k=256 with rescale")

    # Test 3: s_k=256 NOOP rescale (acc_scale=1.0)
    # This should produce the same result as s_k=256 WITHOUT any rescale
    # (which is: O = P[0]*V[0] + P[1]*V[1] with no O rescale — mathematically wrong
    # but should be stable if TMEM round-trip doesn't corrupt)
    out_noop = run_test(256, debug_noop=True, label="s_k=256 NOOP rescale")

    # Test 4: s_k=256 WITHOUT any rescale (old code path)
    # Compare with NOOP to see if TMEM round-trip itself corrupts
    # We can't easily disable the rescale in the current code,
    # but NOOP rescale with factor=1.0 is equivalent to a successful round-trip
    # followed by multiply-by-1. If the output matches a "no-rescale" baseline,
    # the TMEM round-trip is working correctly.


if __name__ == '__main__':
    test()