"""NVFP4-0: Verify Blackwell FP4 primitives are correct.

Quick diagnostic — prints sf_dtype, sf_vec_size, MMA kind, TMA element type.
Run on B200 only.
"""
import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.nvgpu import tcgen05
from cutlass import Float32, BFloat16, Float8E4M3FN, Float8E8M0FNU, Float4E2M1FN
import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))


def test_nvfp4_primitives():
    print("=" * 60)
    print("NVFP4-0.1: sf_dtype and sf_vec_size")
    print("=" * 60)

    from dsv4.ops.quantize import SF_VEC_SIZE
    print(f"  quantize.py SF_VEC_SIZE = {SF_VEC_SIZE}")
    assert SF_VEC_SIZE == 16, f"SF_VEC_SIZE should be 16 for NVFP4, got {SF_VEC_SIZE}"
    print(f"  ✅ SF_VEC_SIZE = 16 (NVFP4 correct)")

    # Construct a BlockScaledGEMM with E4M3 scales (NVFP4)
    M, N, K = 128, 256, 512
    a_fp4 = torch.randint(0, 256, (M, K // 2), device='cuda', dtype=torch.uint8).view(torch.float4_e2m1fn_x2)
    b_fp4 = torch.randint(0, 256, (K // 2, N), device='cuda', dtype=torch.uint8).view(torch.float4_e2m1fn_x2)
    a_sf = torch.randint(0, 256, (M, K // 16), device='cuda', dtype=torch.uint8).view(torch.float8_e4m3fn)
    b_sf = torch.randint(0, 256, (N, K // 16), device='cuda', dtype=torch.uint8).view(torch.float8_e4m3fn)

    from dsv4.kernels.gemm.dense import Sm100BlockScaledPersistentDenseGemmKernel
    try:
        gemm = Sm100BlockScaledPersistentDenseGemmKernel(
            a_ptr=a_fp4,
            b_ptr=b_fp4,
            sfa_ptr=a_sf,
            sfb_ptr=b_sf,
            sf_vec_size=16,
        )
        print(f"  DenseGEMM.sf_vec_size = {gemm.sf_vec_size}")
        print(f"  DenseGEMM.sf_dtype = {gemm.sf_dtype}")
        if gemm.sf_dtype == Float8E4M3FN:
            print(f"  ✅ sf_dtype is Float8E4M3FN (NVFP4 correct)")
        elif gemm.sf_dtype == Float8E8M0FNU:
            print(f"  ⚠️ sf_dtype is Float8E8M0FNU — this is MXFP4 scale format, NOT NVFP4!")
        else:
            print(f"  ❌ sf_dtype is {gemm.sf_dtype} — unexpected!")
    except Exception as e:
        print(f"  DenseGEMM construction failed: {e}")

    # Also try with E8M0 scales (MXFP4)
    a_sf_e8m0 = torch.randint(0, 256, (M, K // 32), device='cuda', dtype=torch.uint8).view(torch.float8_e8m0fnu)
    b_sf_e8m0 = torch.randint(0, 256, (N, K // 32), device='cuda', dtype=torch.uint8).view(torch.float8_e8m0fnu)
    try:
        gemm_e8m0 = Sm100BlockScaledPersistentDenseGemmKernel(
            a_ptr=a_fp4,
            b_ptr=b_fp4,
            sfa_ptr=a_sf_e8m0,
            sfb_ptr=b_sf_e8m0,
            sf_vec_size=32,
        )
        print(f"  DenseGEMM (E8M0/vs=32).sf_dtype = {gemm_e8m0.sf_dtype}")
        print(f"  DenseGEMM (E8M0/vs=32).sf_vec_size = {gemm_e8m0.sf_vec_size}")
    except Exception as e:
        print(f"  DenseGEMM (E8M0) failed: {e}")

    # Check the grouped kernel too
    from dsv4.kernels.gemm.grouped import ScaledGroupedGemmKernel
    try:
        grp = ScaledGroupedGemmKernel(
            a_ptr=a_fp4,
            b_ptr=b_fp4,
            sfa_ptr=a_sf,
            sfb_ptr=b_sf,
            sf_vec_size=16,
        )
        print(f"  GroupedGEMM.sf_vec_size = {grp.sf_vec_size}")
        print(f"  GroupedGEMM.sf_dtype = {grp.sf_dtype}")
    except Exception as e:
        print(f"  GroupedGEMM construction failed: {e}")

    print()
    print("=" * 60)
    print("NVFP4-0.3: FP4 TMA element type in quantize.py")
    print("=" * 60)

    from dsv4.ops.quantize import quantize_tensor_nvfp4
    x = torch.randn(4, 512, device='cuda', dtype=torch.bfloat16)
    x_fp4, x_sf = quantize_tensor_nvfp4(x)
    print(f"  Input dtype: {x.dtype}")
    print(f"  FP4 output dtype: {x_fp4.dtype}")
    print(f"  SF output dtype: {x_sf.dtype}")
    print(f"  FP4 shape: {x_fp4.shape} (expected: [4, 256])")
    print(f"  SF shape: {x_sf.shape} (expected: [4, 32])")
    if x_fp4.dtype == torch.float4_e2m1fn_x2:
        print(f"  ✅ FP4 tensor is float4_e2m1fn_x2 — correct for TMA")
    else:
        print(f"  ❌ FP4 tensor dtype is {x_fp4.dtype} — should be float4_e2m1fn_x2!")

    print()
    print("=" * 60)
    print("NVFP4-0.4: MMA kind verification")
    print("=" * 60)

    try:
        a_major = cutlass.utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
        b_major = cutlass.utils.LayoutEnum.COLUMN_MAJOR.mma_major_mode()
        mma = cutlass.utils.sm100.make_trivial_tiled_mma(
            Float4E2M1FN, Float4E2M1FN, a_major, b_major, Float32,
            tcgen05.CtaGroup.ONE, (128, 256),
            tcgen05.OperandSource.SMEM,
        )
        print(f"  FP4 MMA shape_mnk = {mma.shape_mnk}")
        print(f"  ✅ FP4 MMA construction succeeded")
    except Exception as e:
        print(f"  FP4 MMA construction failed: {e}")

    print()
    print("DONE — NVFP4-0 verification complete")


if __name__ == "__main__":
    test_nvfp4_primitives()