#!/usr/bin/env python3
"""Cosine verification test for B1 mixed FP8/BF16 decode FMHA.

Generates synthetic Q and KV in DSV4 storage format (FP8 noPE + BF16 RoPE),
runs the mixed FP8 decode kernel and the BF16 reference, and compares
per-head cosine similarity.

Production sizes: HD=512, NOPE=448, ROPE=64, N=128..2048, H=128.
"""
import sys
import math
import torch
import torch.nn.functional as F

def quantize_fp8_e4m3(x_fp32):
    """Quantize FP32 tensor to FP8_E4M3 with per-row scale."""
    # x_fp32: (rows, cols)
    amax = x_fp32.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
    scale = amax / 448.0  # E4M3 max representable
    scaled = x_fp32 / scale
    fp8 = scaled.to(torch.float8_e4m3fn)
    return fp8.view(torch.uint8), scale.squeeze(-1)


def run_mixed_fp8_decode(q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=64):
    """Run the B1 mixed FP8 decode FMHA kernel."""
    from dsv4.kernels.attention.fmha_mixed_fp8_op import fmha_mixed_fp8_decode_raw

    B, H, T, HD = q_bf16.shape
    q4 = q_bf16.permute(0, 2, 1, 3).contiguous()  # (B, T, H, HD) -> need (B, H, T, HD)
    q4 = q_bf16  # already (B, H, T, HD)

    o, lse = fmha_mixed_fp8_decode_raw(
        q4, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=rope_dim)
    return o  # (B, H, T, HD) BF16


def run_bf16_reference(q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=64):
    """Run BF16 reference FMHA using PyTorch SDPA on dequantized KV."""
    B, H, T, HD = q_bf16.shape
    NOPE = HD - rope_dim
    N = k_nope_fp8.shape[0]

    # Dequantize FP8 noPE → BF16
    k_nope_flat = k_nope_fp8.view(torch.float8_e4m3fn)
    k_nope_bf16 = k_nope_flat.bfloat16()  # (N, NOPE)
    # Apply per-row scale
    k_nope_bf16 = k_nope_bf16 * k_nope_scale.unsqueeze(-1).bfloat16()

    # Concat noPE + RoPE into full KV
    k_full = torch.cat([k_nope_bf16, k_rope_bf16], dim=-1)  # (N, HD)

    # V = K for MQA (self-attention decode)
    v_full = k_full.clone()

    # Run PyTorch SDPA as reference — FP32 math, exact result
    # q: (B, H, 1, HD), k: (1, 1, N, HD), v: (1, 1, N, HD)
    q_f = q_bf16.float()
    k_f = k_full.float().unsqueeze(0).unsqueeze(0)  # (1, 1, N, HD)
    v_f = v_full.float().unsqueeze(0).unsqueeze(0)  # (1, 1, N, HD)
    # Expand k, v for all batches
    if B > 1:
        k_f = k_f.expand(B, -1, -1, -1)
        v_f = v_f.expand(B, -1, -1, -1)
    o = F.scaled_dot_product_attention(q_f, k_f, v_f, scale=scale)  # (B, H, 1, HD)
    return o.bfloat16()


def test_cosine(N_values, H=128, HD=512, rope_dim=64, B=1, seed=42):
    """Test cosine similarity between mixed FP8 and BF16 reference FMHA."""
    torch.manual_seed(seed)
    NOPE = HD - rope_dim
    scale = 1.0 / math.sqrt(HD)

    all_pass = True
    for N in N_values:
        print(f"\n--- N={N} H={H} HD={HD} ---")

        # Generate synthetic Q (BF16)
        q_fp32 = torch.randn(B, H, 1, HD, dtype=torch.float32) * 0.5
        q_bf16 = q_fp32.bfloat16().cuda()

        # Generate synthetic KV — split into noPE (FP8) + RoPE (BF16)
        k_fp32 = torch.randn(N, HD, dtype=torch.float32) * 0.5
        k_nope_fp32 = k_fp32[:, :NOPE].contiguous()
        k_rope_fp32 = k_fp32[:, NOPE:].contiguous()

        # Quantize noPE to FP8
        k_nope_fp8, k_nope_scale = quantize_fp8_e4m3(k_nope_fp32)
        k_nope_fp8 = k_nope_fp8.cuda()
        k_nope_scale = k_nope_scale.cuda()

        # RoPE stays BF16
        k_rope_bf16 = k_rope_fp32.bfloat16().cuda()

        # Run mixed FP8 decode
        try:
            o_mixed = run_mixed_fp8_decode(q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim)
        except Exception as e:
            print(f"  MIXED FP8 FAILED: {e}")
            all_pass = False
            continue

        # Run BF16 reference
        try:
            o_ref = run_bf16_reference(q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim)
        except Exception as e:
            print(f"  BF16 REF FAILED: {e}")
            all_pass = False
            continue

        # Compare
        o_mixed_f = o_mixed.float()
        o_ref_f = o_ref.float()

        # Global cosine
        cos_global = F.cosine_similarity(o_mixed_f.flatten(), o_ref_f.flatten(), dim=0).item()

        # Per-head cosine (averaged)
        # o shape: (B, H, 1, HD) -> per-head: (B, H, HD)
        o_mixed_h = o_mixed_f.squeeze(2)  # (B, H, HD)
        o_ref_h = o_ref_f.squeeze(2)
        per_head_cos = F.cosine_similarity(o_mixed_h, o_ref_h, dim=-1)  # (B, H)
        min_cos = per_head_cos.min().item()
        mean_cos = per_head_cos.mean().item()

        # Magnitude check
        mixed_max = o_mixed_f.abs().max().item()
        ref_max = o_ref_f.abs().max().item()

        pass_threshold = 0.999
        passed = cos_global >= pass_threshold
        status = "PASS" if passed else "FAIL"

        print(f"  {status}: cos_global={cos_global:.6f} min_head_cos={min_cos:.6f} "
              f"mean_head_cos={mean_cos:.6f}")
        print(f"  |mixed|={mixed_max:.4f} |ref|={ref_max:.4f} "
              f"ratio={mixed_max/ref_max:.4f}" if ref_max > 0 else "  |ref|=0")

        if not passed:
            all_pass = False
            # Print worst heads
            worst = per_head_cos[0].argsort()[:5]
            print(f"  Worst heads: {worst.tolist()} cos={per_head_cos[0][worst].tolist()}")

    return all_pass


if __name__ == "__main__":
    # Production-scale test: N from 128 to 2048
    N_values = [128, 256, 512, 1024, 2048]
    if len(sys.argv) > 1:
        N_values = [int(x) for x in sys.argv[1].split(',')]

    print("=" * 70)
    print("B1 Mixed FP8/BF16 FMHA Cosine Test")
    print("Production sizes: HD=512, NOPE=448, ROPE=64, H=128")
    print("=" * 70)

    all_pass = test_cosine(N_values)

    print("\n" + "=" * 70)
    if all_pass:
        print("ALL TESTS PASSED")
    else:
        print("SOME TESTS FAILED")
        sys.exit(1)
