import itertools
from functools import partial

import pytest
import torch
import triton
import triton.language as tl
from triton_kernels.numerics_details.mxfp import (
    MXFP_BLOCK_SIZE,
    NVFP_BLOCK_SIZE,
    DequantScaleRoundingMode,
    downcast_to_mxfp,
    downcast_to_mxfp_torch,
    get_max_quant_val,
    upcast_from_mxfp,
    upcast_from_mxfp_torch,
)
from triton_kernels.numerics_details.mxfp_details._upcast_from_mxfp import upcast_mxfp4_tile
from triton_kernels.target_info import is_cuda
from triton_kernels.tensor import convert_layout, wrap_torch_tensor
from triton_kernels.tensor_details.layout import StridedLayout
from triton_kernels.testing import assert_close, assert_equal


def dtype_str_to_torch(dtype_str: str) -> torch.dtype:
    return torch.uint8 if dtype_str == "float4_e2m1" else getattr(torch, dtype_str)


@triton.jit
def _upcast_mxfp4_tile_kernel(
    out,
    tensor,
    scale,
    stride_out_m,
    stride_out_k,
    stride_tensor_m,
    stride_tensor_k,
    stride_scale_m,
    stride_scale_k,
    BLOCK_M: tl.constexpr,
    PACKED_K: tl.constexpr,
    SCALE_K: tl.constexpr,
    dst_dtype: tl.constexpr,
):
    offs_m = tl.arange(0, BLOCK_M)[:, None]
    offs_packed_k = tl.arange(0, PACKED_K)[None, :]
    tensor = tl.load(tensor + offs_m * stride_tensor_m + offs_packed_k * stride_tensor_k)

    offs_scale_k = tl.arange(0, SCALE_K)[None, :]
    scale = tl.load(scale + offs_m * stride_scale_m + offs_scale_k * stride_scale_k)

    out_tile = upcast_mxfp4_tile(tensor, scale, dst_dtype)
    offs_k = tl.arange(0, PACKED_K * 2)[None, :]
    tl.store(out + offs_m * stride_out_m + offs_k * stride_out_k, out_tile)


@pytest.mark.skipif(not is_cuda(), reason="Only supported on cuda")
@pytest.mark.parametrize("dst_dtype", ["float16", "bfloat16", "float32"])
def test_mxfp4_tile_upcast_matches_reference(dst_dtype, device):
    torch.manual_seed(0)
    torch.cuda.manual_seed(0)
    dst_dtype = dtype_str_to_torch(dst_dtype)
    rows, k = 64, 128
    block_scales = torch.tensor([0.125, 1.0, 8.0, 64.0], dtype=dst_dtype, device=device)
    block_scales = block_scales.repeat_interleave(MXFP_BLOCK_SIZE.value)
    x = torch.randn((rows, k), dtype=dst_dtype, device=device) * block_scales
    tensor, scale = downcast_to_mxfp(x, torch.uint8, axis=-1)

    ref = upcast_from_mxfp(tensor, scale, dst_dtype, axis=-1)
    out = torch.empty_like(ref)
    tile_dtype = {
        torch.float16: tl.float16,
        torch.bfloat16: tl.bfloat16,
        torch.float32: tl.float32,
    }[dst_dtype]
    _upcast_mxfp4_tile_kernel[(1, )](
        out,
        tensor,
        scale,
        *out.stride(),
        *tensor.stride(),
        *scale.stride(),
        rows,
        tensor.shape[-1],
        scale.shape[-1],
        tile_dtype,
    )

    assert_equal(ref, out)


@pytest.mark.parametrize("dst_dtype", ["float16", "bfloat16", "float32"])
def test_mxfp4_rounding_cases(dst_dtype, device):
    dst_dtype = dtype_str_to_torch(dst_dtype)
    two_point_five_plus_ulp = {
        torch.bfloat16: 0.251953125,
        torch.float16: 0.250244140625,
        torch.float32: 0.2500000298023223877,
    }[dst_dtype]
    pad_values = [0] * 22
    # Construct an example where scale is 1 (when max value is 6.0, the maximum value of e2m1)
    x = torch.tensor([6, 0, 0.24, 0.25, 0.75, 0.99, 1.2, 1.3, -1.25, two_point_five_plus_ulp] + pad_values,
                     dtype=dst_dtype, device=device).view(1, -1, 1)
    quant, scale = downcast_to_mxfp(x, torch.uint8, axis=1)
    dequant = upcast_from_mxfp(quant, scale, dst_dtype, axis=1)
    # Tie-breaking cases (RTNE):
    # - 0.25 is exactly halfway between 0.0 and 0.5. RTNE selects the even quantized value 0.0
    #   (binary LSB of target is 0). Rounding away from zero would pick 0.5; towards zero also picks 0.0.
    # - 0.75 is halfway between 0.5 and 1.0. RTNE selects the even value 1.0 (LSB 0). Away-from-zero would pick 1.0;
    #   towards-zero would pick 0.5.
    # - -1.25 is halfway between -1.0 and -1.5. RTNE selects -1.0 (even). Away-from-zero would pick -1.5;
    #   towards-zero would pick -1.0.
    # - two_point_five_plus_ulp is slightly bigger than 0.25, so it rounds to 0.5.
    assert dequant.flatten().tolist() == [6, 0, 0, 0.0, 1.0, 1.0, 1.0, 1.5, -1.0, 0.5] + pad_values, f"{dequant=}"

    quant_torch, scale_torch = downcast_to_mxfp_torch(x, torch.uint8, axis=1)
    assert_equal(quant_torch, quant)
    assert_equal(scale_torch, scale)

    dequant_torch = upcast_from_mxfp_torch(quant_torch, scale_torch, dst_dtype, axis=1)
    assert_equal(dequant_torch, dequant)

    # ROUND_DOWN should use the max power-of-two when computing scale.
    # Choose a block whose max is 33 so the chosen scale is
    # 2**floor(log2(33/(e2m1 max power of 2 = 4)) = 2**3 = 8 (exponent 127+3),
    # and the other values are multiples of representable FP4 values times 8
    # that allow exact reconstruction.
    pad_values = [0] * 24
    x = torch.tensor([33.0, 24.0, 16.0, 8.0, 4.0, 0.0, -32.0, 0.0] + pad_values,
                     device=device).bfloat16().view(1, -1, 1)
    quant, scale = downcast_to_mxfp(
        x,
        torch.uint8,
        axis=1,
        DEQUANT_SCALE_ROUNDING_MODE=DequantScaleRoundingMode.ROUND_DOWN,
    )
    dequant = upcast_from_mxfp(quant, scale, dst_dtype, axis=1)
    assert_equal(dequant[0, 1:, :], x[0, 1:, :])

    # Golden: scale exponent is 127 + 3 for 2**3 = 8
    assert scale.item() == 127 + 3

    # Torch reference path should match
    quant_torch, scale_torch = downcast_to_mxfp_torch(
        x,
        torch.uint8,
        axis=1,
        DEQUANT_SCALE_ROUNDING_MODE=DequantScaleRoundingMode.ROUND_DOWN,
    )
    assert_equal(quant_torch, quant)
    assert_equal(scale_torch, scale)


@pytest.mark.parametrize("src_dtype", ["float4_e2m1", "float8_e5m2", "float8_e4m3fn"])
@pytest.mark.parametrize("dst_dtype", ["float16", "bfloat16", "float32"])
def test_mxfp_extreme_values(src_dtype, dst_dtype, device):
    if "float8" in src_dtype and (is_cuda() and torch.cuda.get_device_capability()[0] < 9):
        pytest.skip("Float8 not tested on A100")
    src_dtype = dtype_str_to_torch(src_dtype)
    dst_dtype = dtype_str_to_torch(dst_dtype)
    BIG_VALUE = 65470 if dst_dtype == torch.float16 else 3.3895e38
    pad_values = [0] * 30
    x = torch.tensor([BIG_VALUE, BIG_VALUE] + pad_values, dtype=dst_dtype, device=device)
    xq_value, xq_scale = downcast_to_mxfp(x, src_dtype, axis=-1)
    xdq = upcast_from_mxfp(xq_value, xq_scale, dst_dtype, axis=-1)
    xdq_ref = upcast_from_mxfp_torch(xq_value, xq_scale, dst_dtype, axis=-1)
    assert_equal(xdq_ref, xdq)
    assert not xdq.isinf().any()


@pytest.mark.parametrize("src_dtype", ["float4_e2m1", "float8_e5m2", "float8_e4m3fn"])
@pytest.mark.parametrize("dst_dtype", ["float16", "bfloat16", "float32"])
def test_mxfp_quant_dequant(src_dtype, dst_dtype, device):
    if "float8" in src_dtype and (is_cuda() and torch.cuda.get_device_capability()[0] < 9):
        pytest.skip("Float8 not tested on A100")
    limit_range = src_dtype == "float8_e5m2" and dst_dtype == "float16"

    # This test checks that quantization and dequantization kernels produce the exact values for some inputs
    # that can be represented exactly in the quantized format.
    src_dtype = dtype_str_to_torch(src_dtype)
    dst_dtype = dtype_str_to_torch(dst_dtype)
    max_val = get_max_quant_val(src_dtype)
    if limit_range:
        # FP16 can't represent the full range of MXFP8, so we limit the max value here
        max_val = 128

    # These are all the valid mxfp4 positive values.
    pos_vals = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, max_val], device=device, dtype=dst_dtype)
    neg_vals = -pos_vals
    k_dim = torch.cat([pos_vals, neg_vals])
    k_dim = k_dim.reshape([k_dim.shape[0], 1])

    # We pick power of 2 scales since both the scales and their inverse only require exponent bits to be exactly
    # represented. This means we can store the scales exactly in the e8m0 format.
    powers = torch.arange(-8, 8, device=device, dtype=dst_dtype)
    scales = 2**powers
    scales = scales.reshape([1, powers.shape[0]])
    weight = k_dim * scales
    weight = weight.repeat((9, 32))  # Repeat the dimensions to test multi block launches.
    weight = weight.reshape([1, weight.shape[0], weight.shape[1]])
    weight = weight.mT.contiguous().mT
    weight = torch.nn.functional.pad(weight, (0, 0, 0, 16))
    quant, scale = downcast_to_mxfp(weight, src_dtype, axis=1)
    dequant = upcast_from_mxfp(quant, scale, dst_dtype, axis=1)
    assert_equal(weight, dequant)


def test_downcast_to_mxfp_accepts_pitched_strided_input(device):
    torch.manual_seed(0)
    dense = torch.randn((64, 128), device=device, dtype=torch.bfloat16)
    pitched = torch.empty_strided(dense.shape, (256, 1), device=device, dtype=dense.dtype)
    pitched.copy_(dense)

    tensor = convert_layout(wrap_torch_tensor(pitched), StridedLayout(-1))
    quant, scale = downcast_to_mxfp(tensor, torch.uint8, axis=-1)
    expected_quant, expected_scale = downcast_to_mxfp(dense, torch.uint8, axis=-1)

    assert tensor.storage.data.stride() == (256, 1)
    assert_equal(expected_quant, quant)
    assert_equal(expected_scale, scale)


# fmt: off
@pytest.mark.parametrize(
    "shape, axis, quant_dtype, rounding_mode, scale_dtype, microblock_size",
    [
        # Zero-sized arrays
        ((0, 4096, 1024), 1, "float4_e2m1", DequantScaleRoundingMode.ROUND_UP, torch.uint8, MXFP_BLOCK_SIZE.value),
        ((3, 4096, 0), 1, "float4_e2m1", DequantScaleRoundingMode.ROUND_DOWN, torch.uint8, MXFP_BLOCK_SIZE.value),
        ((10, 0, 1024), 2, "float8_e5m2", DequantScaleRoundingMode.ROUND_UP, torch.uint8, MXFP_BLOCK_SIZE.value),
        ((0, 0, 1024), 2, "float8_e4m3fn", DequantScaleRoundingMode.ROUND_DOWN, torch.uint8, MXFP_BLOCK_SIZE.value),
        ((3, 1024, 0), 2, "float4_e2m1", DequantScaleRoundingMode.ROUND_UP, torch.float8_e4m3fn, NVFP_BLOCK_SIZE.value),

        ((3, 4096, 1024), 1, "float4_e2m1", DequantScaleRoundingMode.ROUND_UP, torch.uint8, MXFP_BLOCK_SIZE.value),
        ((32, 254, 60), 0, "float4_e2m1", DequantScaleRoundingMode.ROUND_DOWN, torch.uint8, MXFP_BLOCK_SIZE.value),
        ((1, 320, 160), 2, "float8_e5m2", DequantScaleRoundingMode.ROUND_UP, torch.uint8, MXFP_BLOCK_SIZE.value),
        ((2, 16, 512), -1, "float8_e4m3fn", DequantScaleRoundingMode.ROUND_DOWN, torch.uint8, MXFP_BLOCK_SIZE.value),
        ((2, 64, 3), 1, "float4_e2m1", DequantScaleRoundingMode.ROUND_UP, torch.float8_e4m3fn, NVFP_BLOCK_SIZE.value),
    ],
)
# fmt: on
@pytest.mark.parametrize("dequant_dtype", ["float16", "bfloat16", "float32"])
def test_mxfp_casting(
    shape: tuple[int, ...],
    axis: int,
    quant_dtype: str,
    dequant_dtype: str,
    rounding_mode: DequantScaleRoundingMode,
    scale_dtype: torch.dtype,
    microblock_size: int,
    device,
):
    if ("float8" in quant_dtype
            or scale_dtype == torch.float8_e4m3fn) and (is_cuda() and torch.cuda.get_device_capability()[0] < 9):
        pytest.skip("Float8 not tested on A100")
    torch.manual_seed(0)
    quant_torch_type = dtype_str_to_torch(quant_dtype)
    dequant_torch_type = dtype_str_to_torch(dequant_dtype)
    # Generate random input tensor that is contiguous once axis is the last dimension
    x = torch.randn(shape, device=device, dtype=dequant_torch_type)

    # Quantize and check equivalence
    quant, scale = downcast_to_mxfp(
        x,
        quant_torch_type,
        axis,
        scale_dtype=scale_dtype,
        microblock_size=microblock_size,
        DEQUANT_SCALE_ROUNDING_MODE=rounding_mode,
    )
    quant_torch, scale_torch = downcast_to_mxfp_torch(
        x,
        quant_torch_type,
        axis,
        scale_dtype=scale_dtype,
        microblock_size=microblock_size,
        DEQUANT_SCALE_ROUNDING_MODE=rounding_mode,
    )

    assert_equal(quant_torch, quant)
    assert_equal(scale_torch, scale)
    assert_equal(1, quant.stride(axis))
    assert_equal(1, quant_torch.stride(axis))

    # Dequantize and check equivalence
    dequant = upcast_from_mxfp(quant, scale, dequant_torch_type, axis)
    dequant_torch = upcast_from_mxfp_torch(quant_torch, scale_torch, dequant_torch_type, axis)
    assert_equal(dequant, dequant_torch)

    # Dequantized result should be close to the original, though tolerance is large due to the precision loss.
    assert_close(x, dequant, maxtol=0.5, rmstol=0.15)


def _benchmark_mxfp_quantization(shape, src_dtype: torch.dtype, target_quant_dtype: torch.dtype, n_iters=1000):
    x = torch.randn(*shape, dtype=src_dtype, device="cuda")
    elapsed = (triton.testing.do_bench(
        partial(downcast_to_mxfp, x, target_quant_dtype, axis=-1),
        rep=n_iters,
        return_mode="min",
    ) / 1e3)

    # Each call reads x (2 Bytes) and writes the output tensor (1B or 0.5B) once.
    # -> 3B * numel
    gbytes = ((3 if target_quant_dtype == torch.float8_e4m3fn else 2.5) * x.numel()) / 1e9

    bw = gbytes / elapsed
    return bw


def _benchmark_mxfp_dequantization(shape, src_quant_dtype: torch.dtype, target_dtype: torch.dtype, n_iters=1000):
    x = torch.randn(*shape, dtype=torch.bfloat16, device="cuda").to(src_quant_dtype)
    scale_shape = shape[:-1] + (triton.cdiv(shape[-1], MXFP_BLOCK_SIZE), )
    x_scale = torch.randint(0, 256, scale_shape, device="cuda", dtype=torch.uint8)
    elapsed = (triton.testing.do_bench(
        partial(upcast_from_mxfp, x, x_scale, target_dtype, axis=-1),
        rep=n_iters,
        return_mode="min",
    ) / 1e3)

    # Each call reads x (1B or 0.5B) and writes the output tensor (2 Bytes) once.
    # -> 3B * numel
    gbytes = ((3 if src_quant_dtype == torch.float8_e4m3fn else 2.5) * x.numel()) / 1e9

    bw = gbytes / elapsed
    return bw


if __name__ == "__main__":
    tests = [
        ((1024, 8192), torch.float16),
        ((4096, 8192), torch.float16),
        ((1024, 8192), torch.bfloat16),
        ((4096, 8192), torch.bfloat16),
    ]

    table = []
    shapes = [(1024, 8192), (4096, 8192)]
    source_dtypes = [torch.bfloat16, torch.float16]
    for shape, quant_dtype in itertools.product(shapes, [torch.float8_e4m3fn, torch.uint8]):
        results = [*shape, quant_dtype]
        for src_dtype in source_dtypes:
            results.append(_benchmark_mxfp_quantization(shape, src_dtype, quant_dtype))
        for src_dtype in source_dtypes:
            results.append(_benchmark_mxfp_dequantization(shape, quant_dtype, src_dtype))
        table.append(results)

    from tabulate import tabulate

    headers = [
        "M",
        "N",
        "quant_dtype",
        "quant_bw_bfloat16",
        "quant_bw_float16",
        "dequant_bw_bfloat16",
        "dequant_bw_float16",
    ]
    mxfp8_rows = [row for row in table if row[2] == torch.float8_e4m3fn]
    mxfp4_rows = [row for row in table if row[2] == torch.uint8]

    print("MXFP8 (e4m3fn):")
    print(tabulate(mxfp8_rows, headers=headers))
    print()
    print("MXFP4 (e2m1):")
    print(tabulate(mxfp4_rows, headers=headers))
