"""
Functional tests for MXFP4 / FP4 E2M1 support on AMD gfx950 (CDNA4 / MI350).

All tests are guarded by @tilelang.testing.requires_gfx950 and are silently
skipped on non-gfx950 AMD targets (gfx90a, gfx942, RDNA) and on NVIDIA GPUs.

Test coverage:
  1. FP4 copy operations (global -> shared -> local, cross-type)
  2. Vectorized FP4 <-> float16 / float32 / bfloat16 casts
  3. MXFP4 dequantize-GEMM (fast twiddling + simple path)
"""

import pytest
import torch
import tilelang
import tilelang.testing
import tilelang.language as T
from tilelang import tvm as tvm


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _hip_target():
    return "hip -mcpu=gfx950"


def _fp4_encode(vals):
    """Encode a list of floats to packed uint8 FP4 E2M1 (2 per byte)."""
    fp4_values = [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]

    def encode_one(v):
        best_idx, best_diff = 0, abs(v - fp4_values[0])
        for idx, fv in enumerate(fp4_values):
            diff = abs(v - fv)
            if diff < best_diff:
                best_diff, best_idx = diff, idx
        return best_idx

    nibbles = [encode_one(v) for v in vals]
    assert len(nibbles) % 2 == 0
    return bytes([(nibbles[i] & 0xF) | ((nibbles[i + 1] & 0xF) << 4) for i in range(0, len(nibbles), 2)])


def _fp4_decode(packed: bytes):
    """Decode packed uint8 FP4 E2M1 to list of floats."""
    lut = [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]
    result = []
    for byte in packed:
        result.append(lut[byte & 0xF])
        result.append(lut[(byte >> 4) & 0xF])
    return result


@tilelang.testing.requires_gfx950
def test_static_fp4_shared_allocation_uses_packed_extent():
    @T.prim_func
    def kernel():
        with T.Kernel(1, threads=1):
            A_shared = T.alloc_shared((127,), T.float4_e2m1fn, scope="shared")
            A_shared[126] = T.cast(0.0, T.float4_e2m1fn)

    artifact = tilelang.lower(kernel, target={"kind": "hip", "mcpu": "gfx950"})

    assert "fp4_e2_t A_shared[64];" in artifact.kernel_source


# ---------------------------------------------------------------------------
# Test 1: FP4 copy — shared-memory round-trip
# ---------------------------------------------------------------------------


@tilelang.testing.requires_gfx950
def test_fp4_copy_shared_roundtrip():
    """FP4 values survive a global->shared->global round-trip without corruption."""
    N = 128  # number of FP4 elements (64 bytes packed)
    QN = N // 2

    # Use out_idx=[-1] so dst is allocated by the JIT wrapper.
    @tilelang.jit(out_idx=[-1], target=_hip_target())
    def copy_kernel(N, QN):
        @T.prim_func
        def main(
            src: T.Tensor((QN,), T.uint8),
            dst: T.Tensor((QN,), T.uint8),
        ):
            with T.Kernel(1, threads=64):
                src_sh = T.alloc_shared((QN,), T.uint8)
                T.copy(src[0:QN], src_sh)
                T.copy(src_sh, dst[0:QN])

        return main

    kernel = copy_kernel(N, QN)

    vals = [1.0, -1.5, 2.0, 0.5, 3.0, -3.0, 0.0, 1.0] * (N // 8)
    packed_bytes = _fp4_encode(vals)
    packed = torch.tensor(list(packed_bytes), dtype=torch.uint8).cuda()
    # out_idx=[-1] means dst is allocated internally; pass only src.
    out = kernel(packed)
    assert torch.all(packed == out), "FP4 shared-memory copy corrupted data"


# ---------------------------------------------------------------------------
# Test 2: FP4 -> float16 vectorized cast
# ---------------------------------------------------------------------------


@tilelang.testing.requires_gfx950
def test_fp4_to_float16_cast():
    """FP4 -> float16 cast: dequantize packed uint8 to float16 via the simple path."""
    # The FP4->F16 cast is exercised inside the dequantize GEMM kernel (simple path).
    # This test validates that the simple dequantize path compiles and produces
    # correct results for a small problem size.
    import sys
    import os

    _examples_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "examples", "dequantize_gemm")
    sys.path.insert(0, _examples_dir)
    from example_dequant_gemm_bf16_mxfp4_cdna4 import matmul, ref_program_simple

    M, N, K = 64, 64, 64
    kernel = matmul(
        M,
        N,
        K,
        T.bfloat16,
        T.bfloat16,
        T.float32,
        num_bits=4,
        scale_size=32,
        block_M=64,
        block_N=64,
        block_K=64,
        num_stages=0,
        threads=128,
        split=1,
        fast_dequant=False,
        with_bias=False,
    )
    profiler = kernel.get_profiler(tilelang.TensorSupplyType.Auto)
    profiler.assert_allclose(ref_program_simple, rtol=0.02, atol=0.02)


# ---------------------------------------------------------------------------
# Test 3: MXFP4 dequantize GEMM - simple path
# ---------------------------------------------------------------------------


@tilelang.testing.requires_gfx950
@pytest.mark.parametrize("M,N,K", [(256, 256, 256), (128, 512, 128)])
def test_mxfp4_dequant_gemm_simple(M, N, K):
    """MXFP4 dequantize-GEMM (simple path) produces correct BF16 output."""
    import sys
    import os

    _examples_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "examples", "dequantize_gemm")
    sys.path.insert(0, _examples_dir)
    from example_dequant_gemm_bf16_mxfp4_cdna4 import matmul, ref_program_simple

    scale_size = 32
    kernel = matmul(
        M,
        N,
        K,
        T.bfloat16,
        T.bfloat16,
        T.float32,
        num_bits=4,
        scale_size=scale_size,
        block_M=128,
        block_N=128,
        block_K=128,
        num_stages=2,
        threads=256,
        split=1,
        fast_dequant=False,
        with_bias=False,
    )
    profiler = kernel.get_profiler(tilelang.TensorSupplyType.Auto)
    profiler.assert_allclose(ref_program_simple, rtol=0.02, atol=0.02)


# ---------------------------------------------------------------------------
# Test 4: MXFP4 dequantize GEMM - fast twiddling path
# ---------------------------------------------------------------------------


@tilelang.testing.requires_gfx950
@pytest.mark.parametrize("M,N,K", [(256, 256, 256)])
def test_mxfp4_dequant_gemm_twiddling(M, N, K):
    """MXFP4 dequantize-GEMM (fast twiddling path) produces correct BF16 output."""
    import sys
    import os

    _examples_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "examples", "dequantize_gemm")
    sys.path.insert(0, _examples_dir)
    from example_dequant_gemm_bf16_mxfp4_cdna4 import matmul, ref_program_twiddling

    scale_size = 32
    kernel = matmul(
        M,
        N,
        K,
        T.bfloat16,
        T.bfloat16,
        T.float32,
        num_bits=4,
        scale_size=scale_size,
        block_M=128,
        block_N=128,
        block_K=128,
        num_stages=2,
        threads=256,
        split=1,
        fast_dequant=True,
        with_bias=False,
    )
    profiler = kernel.get_profiler(tilelang.TensorSupplyType.Auto)
    profiler.assert_allclose(ref_program_twiddling, rtol=0.02, atol=0.02)


# ---------------------------------------------------------------------------
# Test 5: get_mxfp_intrin_group returns HIP source for gfx950
# ---------------------------------------------------------------------------


@tilelang.testing.requires_gfx950
def test_get_mxfp_intrin_group_returns_hip_source():
    """get_mxfp_intrin_group() returns HIP C++ source (not CUDA PTX) for gfx950."""
    from tilelang.quantize import get_mxfp_intrin_group
    from tilelang import tvm

    target = tvm.target.Target("hip -mcpu=gfx950")
    info = get_mxfp_intrin_group(
        out_dtype=T.bfloat16,
        source_bit=4,
        use_twiddling=True,
        target=target,
    )
    assert "func_name" in info and "c_source" in info
    # HIP source uses __device__ and does NOT contain PTX asm keywords
    src = info["c_source"]
    assert "__device__" in src, "Expected __device__ in HIP source"
    assert "prmt.b32" not in src, "Expected no PTX asm in HIP source (should be C++)"
    assert "decode_fp4_to_bf16_twiddling" in info["func_name"]


# ---------------------------------------------------------------------------
# Test 6: get_mxfp_intrin_group for non-gfx950 still returns CUDA PTX
# ---------------------------------------------------------------------------


def test_get_mxfp_intrin_group_returns_ptx_for_cuda():
    """get_mxfp_intrin_group() returns CUDA PTX source when target is None."""
    from tilelang.quantize import get_mxfp_intrin_group

    info = get_mxfp_intrin_group(
        out_dtype=T.bfloat16,
        source_bit=4,
        use_twiddling=True,
        target=None,  # default: CUDA/NV path
    )
    src = info["c_source"]
    assert "prmt.b32" in src, "Expected PTX asm in CUDA source"


if __name__ == "__main__":
    tilelang.testing.main()
