#!/usr/bin/env python3
"""
NVFP4 Kernel Test - Validates E2M1/E4M3 encoding/decoding and matmul on GPU.

Tests:
1. E2M1 quantization round-trip (quantize -> dequantize -> compare)
2. E4M3 scale encoding round-trip
3. NVFP4 block quantization + dequantization accuracy
4. NVFP4 matmul vs FP32 reference

Run with: python3 /data-nvme/test_nvfp4.py
Or via nvrunner: python3 /data-nvme/nvrunner python3 /data-nvme/test_nvfp4.py --duration 120
"""

import numpy as np
import torch

# ============================================================================
# E2M1 constants
# ============================================================================

E2M1_VALUES = np.array([
    0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
    -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0
], dtype=np.float32)

NVFP4_BLOCK_SIZE = 16
NVFP4_MAX_VAL = 6.0

# ============================================================================
# E2M1 encoding/decoding (Python reference)
# ============================================================================

def float_to_e2m1(x):
    """Convert float to nearest E2M1 4-bit code."""
    if x == 0.0:
        return 0
    sign = 1 if x < 0 else 0
    abs_x = abs(x)
    best_code = 0
    best_err = float('inf')
    for i in range(8):
        err = abs(abs_x - E2M1_VALUES[i])
        if err < best_err:
            best_err = err
            best_code = i
    return (best_code | (0x8 if sign else 0))

def e2m1_to_float(code):
    """Convert E2M1 4-bit code to float."""
    return E2M1_VALUES[code & 0xF]

def float_to_e4m3(x):
    """Convert float to E4M3 (unsigned, for scale factors)."""
    if x == 0.0:
        return 0
    x = abs(x)
    if x >= 448.0:
        return 0x7F
    
    bits = np.float32(x).view(np.uint32)
    exp = int((bits >> 23) & 0xFF)
    mant = int(bits & 0x7FFFFF)
    
    if exp == 0:
        return 0
    
    real_exp = exp - 127
    e4m3_exp = real_exp + 7
    
    if e4m3_exp <= 0:
        mant_frac = x * 64.0
        mant_int = int(round(mant_frac))
        if mant_int >= 8:
            return 0x08
        return mant_int & 0x07
    
    if e4m3_exp > 15:
        return 0x7F
    
    e4m3_mant = mant >> 20
    round_bit = (mant >> 19) & 1
    sticky = (mant & ((1 << 19) - 1)) != 0
    
    if round_bit and (sticky or (e4m3_mant & 1)):
        e4m3_mant += 1
        if e4m3_mant >= 8:
            e4m3_mant = 0
            e4m3_exp += 1
            if e4m3_exp > 15:
                return 0x7F
    
    return (e4m3_exp << 3) | (e4m3_mant & 0x07)

def e4m3_to_float(code):
    """Convert E4M3 to float."""
    if code == 0:
        return 0.0
    sign = (code >> 7) & 1
    exp = (code >> 3) & 0x0F
    mant = code & 0x07
    
    if exp == 0:
        val = mant * 0.001953125  # 2^(-9)
    else:
        val = (1.0 + mant / 8.0) * (2.0 ** (exp - 7))
    
    return -val if sign else val

# ============================================================================
# NVFP4 block quantization/dequantization (Python reference)
# ============================================================================

def quantize_nvfp4_block(values):
    """Quantize a block of 16 float values to NVFP4.
    Returns (packed_data[8], scale_e4m3).
    """
    assert len(values) == NVFP4_BLOCK_SIZE
    
    # Compute block amax
    block_amax = max(abs(v) for v in values)
    
    # block_scale = block_amax / 6.0
    block_scale = block_amax / NVFP4_MAX_VAL if block_amax > 0 else 0.0
    
    # Convert to E4M3
    scale_e4m3 = float_to_e4m3(block_scale)
    scale_float = e4m3_to_float(scale_e4m3)
    
    # Quantize each value
    packed = [0] * 8
    for i in range(NVFP4_BLOCK_SIZE):
        if scale_float > 0:
            normalized = values[i] / scale_float
        else:
            normalized = 0.0
        code = float_to_e2m1(normalized)
        
        if i % 2 == 0:
            packed[i // 2] = (packed[i // 2] & 0xF0) | (code & 0x0F)
        else:
            packed[i // 2] = (packed[i // 2] & 0x0F) | ((code & 0x0F) << 4)
    
    return packed, scale_e4m3

def dequantize_nvfp4_block(packed, scale_e4m3):
    """Dequantize an NVFP4 block.
    Returns 16 float values.
    """
    scale_float = e4m3_to_float(scale_e4m3)
    values = [0.0] * NVFP4_BLOCK_SIZE
    
    for i in range(NVFP4_BLOCK_SIZE):
        byte = packed[i // 2]
        code = (byte & 0x0F) if (i % 2 == 0) else ((byte >> 4) & 0x0F)
        values[i] = e2m1_to_float(code) * scale_float
    
    return values

# ============================================================================
# GPU kernel test using PyTorch (as proxy for CUDA)
# ============================================================================

def test_e2m1_roundtrip():
    """Test E2M1 encoding/decoding accuracy."""
    print("=== Test 1: E2M1 Round-Trip ===")
    test_values = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
                           -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, 0.25], dtype=np.float32)
    
    max_err = 0.0
    for v in test_values:
        code = float_to_e2m1(v)
        decoded = e2m1_to_float(code)
        err = abs(v - decoded)
        max_err = max(max_err, err)
        print(f"  {v:8.4f} -> code={code:2d} (0x{code:X}) -> {decoded:8.4f}  err={err:.4f}")
    
    assert max_err <= 0.5, f"E2M1 max error {max_err} exceeds 0.5"
    print(f"  PASSED (max_err={max_err:.4f})\n")

def test_e4m3_roundtrip():
    """Test E4M3 scale encoding/decoding."""
    print("=== Test 2: E4M3 Scale Round-Trip ===")
    # E4M3 subnormal min = 2^(-9) = 0.001953125
    # Normal range starts at 2^(-6) = 0.015625
    # Block scales for typical weights: block_amax/6.0, where amax > 0.1
    test_scales = [0.05, 0.1, 0.5, 1.0, 2.0, 10.0, 100.0, 448.0]
    
    max_rel_err = 0.0
    for s in test_scales:
        code = float_to_e4m3(s)
        decoded = e4m3_to_float(code)
        rel_err = abs(s - decoded) / max(s, 1e-10)
        max_rel_err = max(max_rel_err, rel_err)
        print(f"  {s:10.6f} -> code=0x{code:02X} -> {decoded:10.6f}  rel_err={rel_err:.6f}")
    
    # E4M3 has 3 mantissa bits, so worst case relative error for normals is ~6.25%
    # Subnormals have higher error, but we exclude extreme subnormals
    assert max_rel_err < 0.1, f"E4M3 max relative error {max_rel_err} exceeds 0.1"
    print(f"  PASSED (max_rel_err={max_rel_err:.6f})\n")

def test_nvfp4_block_roundtrip():
    """Test NVFP4 block quantization round-trip."""
    print("=== Test 3: NVFP4 Block Round-Trip ===")
    np.random.seed(42)
    
    # Generate test data with various distributions
    test_blocks = [
        np.random.randn(NVFP4_BLOCK_SIZE).astype(np.float32) * 0.1,  # Small values
        np.random.randn(NVFP4_BLOCK_SIZE).astype(np.float32) * 1.0,  # Normal
        np.random.randn(NVFP4_BLOCK_SIZE).astype(np.float32) * 10.0, # Large
        np.linspace(-5, 5, NVFP4_BLOCK_SIZE, dtype=np.float32),       # Uniform
        np.zeros(NVFP4_BLOCK_SIZE, dtype=np.float32),                 # All zeros
    ]
    
    for i, block in enumerate(test_blocks):
        packed, scale = quantize_nvfp4_block(block.tolist())
        dequant = dequantize_nvfp4_block(packed, scale)
        
        max_err = max(abs(a - b) for a, b in zip(block, dequant))
        max_val = max(abs(v) for v in block)
        rel_err = max_err / max(max_val, 1e-10)
        
        print(f"  Block {i}: max_val={max_val:.4f}, max_err={max_err:.6f}, rel_err={rel_err:.6f}")
        
        # NVFP4 has 8 representable positive values, so error should be bounded
        assert rel_err < 0.3, f"Block {i} relative error {rel_err} too high"
    
    print("  PASSED\n")

def test_nvfp4_matmul_gpu():
    """Test NVFP4 quantized matmul on GPU using PyTorch.
    
    This simulates the NVFP4 kernel by doing:
    1. Quantize weight matrix to NVFP4
    2. Dequantize on-the-fly during matmul
    3. Compare with FP32 reference
    """
    print("=== Test 4: NVFP4 MatMul GPU (PyTorch) ===")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"  Device: {device}")
    
    if device.type == "cpu":
        print("  SKIPPED (no CUDA)\n")
        return
    
    # Test dimensions
    M, N, K = 1, 256, 512
    
    # Generate test data
    torch.manual_seed(42)
    x = torch.randn(M, K, device=device, dtype=torch.float32)
    w = torch.randn(N, K, device=device, dtype=torch.float32) * 0.5  # Smaller weights for better quantization
    
    # FP32 reference
    ref = x @ w.t()  # [M, N]
    
    # NVFP4 quantization of weights (on CPU then transfer)
    num_blocks_k = K // NVFP4_BLOCK_SIZE
    w_quant_np = np.zeros((N, K), dtype=np.float32)
    w_np = w.cpu().numpy()
    
    for n in range(N):
        for bk in range(num_blocks_k):
            block = w_np[n, bk*NVFP4_BLOCK_SIZE:(bk+1)*NVFP4_BLOCK_SIZE]
            block_amax = np.abs(block).max()
            block_scale = block_amax / NVFP4_MAX_VAL if block_amax > 0 else 0.0
            
            scale_e4m3 = float_to_e4m3(block_scale)
            scale_float = e4m3_to_float(scale_e4m3)
            
            for i in range(NVFP4_BLOCK_SIZE):
                if scale_float > 0:
                    normalized = block[i] / scale_float
                else:
                    normalized = 0.0
                code = float_to_e2m1(normalized)
                w_quant_np[n, bk*NVFP4_BLOCK_SIZE + i] = e2m1_to_float(code) * scale_float
    
    w_quant = torch.from_numpy(w_quant_np).to(device)
    
    # NVFP4 matmul
    result = x @ w_quant.t()  # [M, N]
    
    # Compare
    max_err = (ref - result).abs().max().item()
    mean_err = (ref - result).abs().mean().item()
    ref_norm = ref.abs().mean().item()
    rel_err = mean_err / max(ref_norm, 1e-10)
    
    print(f"  Dimensions: M={M}, N={N}, K={K}")
    print(f"  Max error: {max_err:.6f}")
    print(f"  Mean error: {mean_err:.6f}")
    print(f"  Relative error: {rel_err:.6f}")
    
    assert rel_err < 0.2, f"MatMul relative error {rel_err} too high"
    print("  PASSED\n")

def test_nvfp4_large_matmul_gpu():
    """Test with larger dimensions to stress the kernel."""
    print("=== Test 5: NVFP4 Large MatMul GPU ===")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    if device.type == "cpu":
        print("  SKIPPED (no CUDA)\n")
        return
    
    M, N, K = 4, 1024, 2048
    
    torch.manual_seed(123)
    x = torch.randn(M, K, device=device, dtype=torch.float32)
    w = torch.randn(N, K, device=device, dtype=torch.float32) * 0.3
    
    ref = x @ w.t()
    
    # Vectorized NVFP4 quantization
    num_blocks_k = K // NVFP4_BLOCK_SIZE
    w_blocks = w.view(N, num_blocks_k, NVFP4_BLOCK_SIZE)
    block_amax = w_blocks.abs().amax(dim=2, keepdim=True)
    block_scale = block_amax / NVFP4_MAX_VAL
    block_scale[block_scale == 0] = 1.0  # Avoid division by zero
    
    # Quantize: round to nearest E2M1 value
    normalized = w_blocks / block_scale
    # Find nearest E2M1 value
    e2m1_vals_t = torch.tensor(E2M1_VALUES, device=device, dtype=torch.float32)
    # For each normalized value, find nearest in E2M1_VALUES
    expanded = normalized.unsqueeze(-1)  # [N, num_blocks, 16, 1]
    diffs = (expanded - e2m1_vals_t).abs()  # [N, num_blocks, 16, 16]
    codes = diffs.argmin(dim=-1)  # [N, num_blocks, 16]
    w_quant = e2m1_vals_t[codes] * block_scale
    w_quant = w_quant.view(N, K)
    
    result = x @ w_quant.t()
    
    max_err = (ref - result).abs().max().item()
    mean_err = (ref - result).abs().mean().item()
    ref_norm = ref.abs().mean().item()
    rel_err = mean_err / max(ref_norm, 1e-10)
    
    print(f"  Dimensions: M={M}, N={N}, K={K}")
    print(f"  Max error: {max_err:.6f}")
    print(f"  Mean error: {mean_err:.6f}")
    print(f"  Relative error: {rel_err:.6f}")
    
    assert rel_err < 0.2, f"Large MatMul relative error {rel_err} too high"
    print("  PASSED\n")

def test_memory_savings():
    """Verify NVFP4 memory savings vs FP32."""
    print("=== Test 6: Memory Savings ===")
    
    K = 4096
    N = 4096
    
    fp32_bytes = N * K * 4
    nvfp4_bytes = N * (K // NVFP4_BLOCK_SIZE) * 9  # 9 bytes per 16 elements
    
    ratio = fp32_bytes / nvfp4_bytes
    
    print(f"  FP32: {fp32_bytes / 1024 / 1024:.1f} MB")
    print(f"  NVFP4: {nvfp4_bytes / 1024 / 1024:.1f} MB")
    print(f"  Compression ratio: {ratio:.2f}x")
    
    # 4 bits per element + 8 bits per 16 elements scale = 4.5 bits per element
    # vs 32 bits per element for FP32
    expected_ratio = 32.0 / 4.5
    assert abs(ratio - expected_ratio) < 0.1, f"Compression ratio {ratio} != expected {expected_ratio}"
    print(f"  PASSED (expected ~{expected_ratio:.1f}x)\n")

if __name__ == "__main__":
    print("NVFP4 Kernel Test Suite\n")
    print(f"GPU available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
    print()
    
    test_e2m1_roundtrip()
    test_e4m3_roundtrip()
    test_nvfp4_block_roundtrip()
    test_nvfp4_matmul_gpu()
    test_nvfp4_large_matmul_gpu()
    test_memory_savings()
    
    print("=== ALL TESTS PASSED ===")