"""
P6 Integration Test: One-way TMEM→regs→SMEM→GMEM epilogue.

Verifies the epilogue refactoring:
  1. TMEM → registers (tcgen05.ld, warp-collective)
  2. epilogue_op in registers (normalize, FP4 hook)
  3. Registers → SMEM (row-major)
  4. SMEM → GMEM (direct write)

Gate: worst-case cosine >= 0.999990 per configuration.
"""
import torch
import math
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def cosine_sim(a, b):
    a = a.flatten().float()
    b = b.flatten().float()
    return (a @ b) / (a.norm() * b.norm() + 1e-30)


def reference_attention(q_4d, k_4d, v_4d, scale):
    """PyTorch reference matching kernel tensor layout."""
    n_h = q_4d.shape[1]
    n_kv = k_4d.shape[1]
    N = k_4d.shape[2]
    q_per_kv = n_h // n_kv
    q = q_4d[0]
    k = k_4d[0]
    v = v_4d[0].transpose(-1, -2)

    output = torch.zeros(n_h, 1, q_4d.shape[3], dtype=torch.bfloat16, device='cuda')
    for h in range(n_h):
        kv_idx = h // q_per_kv
        q_h = q[h]
        k_h = k[kv_idx]
        v_h = v[kv_idx]
        s = torch.matmul(q_h.float(), k_h.float().T) * scale
        s = torch.softmax(s, dim=-1)
        o = torch.matmul(s, v_h.float())
        output[h] = o.bfloat16()
    return output


def test_one_way_epilogue():
    """Test the one-way TMEM→regs→SMEM→GMEM epilogue path."""
    from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw
    from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw

    torch.manual_seed(42)
    configs = [
        # Single-tile (N<=128)
        (4, 4, 64, 64, "MHA hd=64", False),
        (4, 4, 128, 128, "MHA hd=128", False),
        (4, 4, 64, 256, "MHA hd=256", False),
        (4, 1, 64, 64, "MQA hd=64", False),
        (128, 1, 64, 64, "MQA Pro hd=64", False),
        (8, 2, 64, 64, "GQA hd=64", False),
        # Multi-tile (N>128)
        (4, 4, 256, 64, "MHA hd=64 N=256 (2 tiles)", True),
        (4, 4, 512, 64, "MHA hd=64 N=512 (4 tiles)", True),
        (4, 1, 256, 128, "MQA hd=128 N=256 (2 tiles)", True),
    ]

    all_pass = True
    for n_q, n_kv, N, hd, desc, multitile in configs:
        scale = 1.0 / math.sqrt(hd)
        q_4d = torch.randn(1, n_q, 1, hd, dtype=torch.bfloat16, device='cuda').contiguous()
        k_4d = torch.randn(1, n_kv, N, hd, dtype=torch.bfloat16, device='cuda').contiguous()
        v_4d = torch.randn(1, n_kv, hd, N, dtype=torch.bfloat16, device='cuda').contiguous()

        if multitile:
            o_4d, _ = fmha_multitile_decode_raw(q_4d, k_4d, v_4d, scale)
        else:
            sb = torch.zeros(1, n_q, dtype=torch.float32, device='cuda')
            o_4d, _ = fmha_multitile_decode_raw(q_4d, k_4d, v_4d, scale)

        o_ref = reference_attention(q_4d, k_4d, v_4d, scale)

        worst_cos = 1.0
        for h in range(n_q):
            cos = torch.nn.functional.cosine_similarity(
                o_4d[0, h].flatten().float().unsqueeze(0),
                o_ref[h].flatten().float().unsqueeze(0)
            ).item()
            worst_cos = min(worst_cos, cos)

        status = "PASS" if worst_cos >= 0.999994 else "FAIL"
        if worst_cos < 0.999994:
            all_pass = False
        print(f"  {status} {desc}: worst_cos={worst_cos:.6f}")

    return all_pass


if __name__ == "__main__":
    print("P6 Integration Test: One-way TMEM→regs→SMEM→GMEM epilogue")
    print("=" * 60)

    if test_one_way_epilogue():
        print("ALL PASS")
    else:
        print("SOME FAILED")
        sys.exit(1)
