# isort: off
# fmt: off
from enum import Enum
import math
import triton
import torch
import torch.nn.functional as F
from .mxfp_details._upcast_from_mxfp import _upcast_from_mxfp
from .mxfp_details._downcast_to_mxfp import _downcast_to_mxfp, MXFP_BLOCK_SIZE, NVFP_BLOCK_SIZE, _quantize_mxfp8_fn, _quantize_mxfp4_fn, _quantize_nvfp4_fn
from triton.tools.tensor_descriptor import TensorDescriptor
from triton_kernels.tensor import Tensor, wrap_torch_tensor, empty
from triton_kernels.tensor_details.layout import StridedLayout
from triton_kernels.tensor_details.dtype import FP4, FP8_E4M3FN, FP8_E5M2, UINT8
# -----------------------------------------------------------------------------
#                      Dequantization / Quantization Utilities
# -----------------------------------------------------------------------------

class DequantScaleRoundingMode(Enum):
    # 2^round_up(log2(max/max_q)) avoids clipping the max value
    ROUND_UP = 0
    # 2^round_down(log2(max/max_power_of_2_q)) follows the OCP standard ~50% of
    # chance of clipping the max value.
    ROUND_DOWN = 1

def downcast_to_mxfp(x: torch.Tensor, out_dtype: torch.dtype, axis: int,
                     scale_dtype: torch.dtype = torch.uint8,
                     microblock_size: int = MXFP_BLOCK_SIZE.value,
                     DEQUANT_SCALE_ROUNDING_MODE: DequantScaleRoundingMode = DequantScaleRoundingMode.ROUND_UP):
    """
         Convert the src weights to mx format. The src weight is quantized along the axis dimension.

         If weight_quant_type is torch.uint8, we output mxfp4/NVFP4 where two e2m1 values are packed into a single byte.
         Note that this means the k_dim of the tensor will be half of the logical k_dim.

         If weight_quant_type is torch.float8_e4m3fn or torch.float8_e5m2, we output mxfp8 with the float8s are stored
         in their respective formats.
    """
    if not isinstance(x, Tensor):
        x = wrap_torch_tensor(x)
    assert scale_dtype == torch.uint8 or scale_dtype == torch.float8_e4m3fn, f"Invalid scale dtype {scale_dtype=}"
    if isinstance(out_dtype, torch.dtype):
        out_dtype = {
            torch.uint8: FP4,
            torch.float8_e4m3fn: FP8_E4M3FN,
            torch.float8_e5m2: FP8_E5M2,
        }[out_dtype]
    assert x.shape[axis] % microblock_size == 0, f"axis dim must be divisible by {microblock_size}. Got {x.shape[axis]}"
    assert isinstance(x.storage.layout, StridedLayout), "input data must be strided"
    assert -x.ndim <= axis < x.ndim, f"Invalid axis {axis=}"
    assert out_dtype in (FP4, FP8_E4M3FN, FP8_E5M2), f"Invalid output dtype {out_dtype=}"
    if scale_dtype == torch.float8_e4m3fn:
        assert out_dtype == FP4, f"Direct float8 scales are only supported for FP4 values. Got {out_dtype=}"
        assert DEQUANT_SCALE_ROUNDING_MODE == DequantScaleRoundingMode.ROUND_UP, \
            "Direct float8 scales only support ROUND_UP in downcast_to_mxfp"
    # handle negative `axis``
    axis = axis if axis >= 0 else axis + x.ndim
    # downcast
    L = x.shape[axis]
    # Ensure last dimension is a multiple of the microblock size. This is expected by the kernel.
    # output value storage
    y_layout = StridedLayout(major_dim=axis - x.ndim)
    y_scale_shape = (*x.shape[:axis], triton.cdiv(L, microblock_size), *x.shape[axis+1:])
    y_value = empty(x.shape, out_dtype, x.device, y_layout)
    y_scale = empty(y_scale_shape, UINT8 if scale_dtype == torch.uint8 else FP8_E4M3FN, x.device, y_layout)
    if x.numel() > 0:
        # canonicalize to a 2D tensor that paxks 4-bit values on its inner-most dimension
        x_storage = x.storage.data.transpose(axis, -1).reshape(-1, x.shape[axis])
        y_storage_value = y_value.storage.data.transpose(axis, -1).view(-1, y_value.storage.data.shape[axis])
        y_storage_scale = y_scale.storage.data.transpose(axis, -1).view(-1, y_scale.storage.data.shape[axis])
        # performance hyper-parameters
        BLOCK_OUT_DIM = 32
        BLOCK_QUANT_DIM = microblock_size * 4
        NUM_WARPS = 4 if x.dtype == torch.float32 else 8
        # launch kernel
        blocks_out_dim = triton.cdiv(x_storage.shape[0], BLOCK_OUT_DIM)
        blocks_quant_dim = triton.cdiv(x_storage.shape[1], BLOCK_QUANT_DIM)
        _downcast_to_mxfp[(blocks_out_dim, blocks_quant_dim)](
            y_storage_value, *y_storage_value.stride(),
            y_storage_scale, *y_storage_scale.stride(),
            x_storage, *x_storage.stride(), *x_storage.shape,
            BLOCK_OUT_DIM,
            BLOCK_QUANT_DIM,
            microblock_size,
            DEQUANT_SCALE_ROUNDING_MODE.value,
            num_warps=NUM_WARPS,
        )
    # TODO: return tensor object instead of its storage
    return y_value.storage.data, y_scale.storage.data

def upcast_from_mxfp(
    tensor: torch.Tensor,
    scale: torch.Tensor,
    target_dtype: torch.dtype,
    axis: int,
    tensor_scale: torch.Tensor | None = None,
):
    """
    Upcasts an mxfp (packed) weight tensor back to float16 or bfloat16.

    The function assumes that the tensors were quantized along the given axis.
    It permutes the tensor so that the quantized axis is last, reshapes to 2D,
    launches the Triton upcast kernel, and then unpermutes back to the original order.
    """
    ndim = tensor.ndim
    assert -ndim <= axis < ndim, f"Invalid axis {axis=}"
    axis = axis if axis >= 0 else axis + ndim
    assert tensor.ndim == scale.ndim, (f"Weight and scale must have the same number of dimensions. "
                                       f"Got {tensor.ndim=} and {scale.ndim=}")
    if tensor_scale is not None and tensor_scale.numel() != 1:
        expected_tensor_scale_shape = tensor.shape[:axis] + tensor.shape[axis + 1:]
        assert tensor_scale.shape == expected_tensor_scale_shape, \
            f"Invalid tensor scale shape: expected {expected_tensor_scale_shape}, got {tensor_scale.shape}"
    # dtype checks
    assert tensor.dtype in {torch.uint8, torch.float8_e5m2, torch.float8_e4m3fn}, \
        f"Invalid tensor dtype {tensor.dtype=}"
    assert scale.dtype == torch.uint8 or scale.dtype == torch.float8_e4m3fn, \
        f"Invalid scale dtype {scale.dtype=}"
    assert target_dtype in (torch.float16, torch.bfloat16, torch.float32), f"Invalid output dtype {target_dtype=}"
    # upcast
    pack_multiple = 2 if tensor.dtype == torch.uint8 else 1
    logical_quant_dim = tensor.shape[axis] * pack_multiple
    tensor = tensor.transpose(axis, tensor.ndim - 1).contiguous()
    scale = scale.transpose(axis, scale.ndim - 1).contiguous()
    original_out_shape = tensor.shape[:-1] + (logical_quant_dim, )

    if tensor.numel() > 0:
        reshaped_tensor = tensor.view(-1, tensor.shape[-1])
        reshaped_scale = scale.view(-1, scale.shape[-1])

        # Pad the tensor and output if needed for tensor descriptor spec requirements.
        TENSOR_DESC_PAD_REQ = 16
        needs_padding = reshaped_tensor.shape[-1] % TENSOR_DESC_PAD_REQ != 0
        if needs_padding:
            tensor_pad_amount = TENSOR_DESC_PAD_REQ - (reshaped_tensor.shape[-1] % TENSOR_DESC_PAD_REQ)
            reshaped_tensor = F.pad(reshaped_tensor, (0, tensor_pad_amount), "constant", 0)
            pad_elems_count = tensor_pad_amount * pack_multiple
            out_shape = original_out_shape[:-1] + (original_out_shape[-1] + pad_elems_count, )
        else:
            out_shape = original_out_shape
        out = torch.empty(out_shape, dtype=target_dtype, device=tensor.device)
        reshaped_out = out.view(-1, out.shape[-1])

        is_fp4 = reshaped_tensor.dtype == torch.uint8
        scale_block_size = MXFP_BLOCK_SIZE.value if scale.dtype == torch.uint8 else NVFP_BLOCK_SIZE.value

        # performance hyper-parameters
        BLOCK_OUT_DIM = 64
        BLOCK_QUANT_DIM = scale_block_size * 4
        NUM_WARPS = 4

        blocks_out_dim = triton.cdiv(reshaped_out.shape[0], BLOCK_OUT_DIM)
        blocks_quant_dim = triton.cdiv(reshaped_out.shape[1], BLOCK_QUANT_DIM)
        k_divisor = 2 if is_fp4 else 1
        block_size_quant_mx_tensor = BLOCK_QUANT_DIM // k_divisor
        out_desc = TensorDescriptor.from_tensor(reshaped_out, [BLOCK_OUT_DIM, BLOCK_QUANT_DIM])
        tensor_desc = TensorDescriptor.from_tensor(reshaped_tensor, [BLOCK_OUT_DIM, block_size_quant_mx_tensor])
        _upcast_from_mxfp[(blocks_out_dim, blocks_quant_dim)](
            out_desc,
            tensor_desc,
            reshaped_scale,
            *reshaped_scale.stride(),
            *reshaped_out.shape,
            BLOCK_OUT_DIM,
            BLOCK_QUANT_DIM,
            scale_block_size,
            num_warps=NUM_WARPS,
        )
        if needs_padding:
            out = out[..., :original_out_shape[-1]]
    else:
        out = torch.empty(original_out_shape, dtype=target_dtype, device=tensor.device)
    out = out.transpose(axis, scale.ndim - 1).contiguous()
    if tensor_scale is not None:
        tensor_scale_shape = (1,) * out.ndim if tensor_scale.numel() == 1 else (*tensor_scale.shape[:axis], 1, *tensor_scale.shape[axis:])
        out = out * tensor_scale.to(device=out.device).reshape(tensor_scale_shape)
    return out


# ------------


def right_shift_unsigned(x, shift):
    # CUDA torch does not support bit ops on uint32, so we need to mask to get unsigned right shift
    return (x >> shift) & ((1 << (32 - shift)) - 1)


def get_max_quant_val(dtype: torch.dtype):
    d = {torch.uint8: 6.0, torch.float8_e5m2: 57344.0, torch.float8_e4m3fn: 448.0}
    return d[dtype]


def downcast_to_mxfp_torch(src_tensor: torch.Tensor, out_quant_type: torch.dtype, axis: int,
                           scale_dtype: torch.dtype = torch.uint8,
                           microblock_size: int = MXFP_BLOCK_SIZE.value,
                           expected_scale: torch.Tensor | None = None,
                           DEQUANT_SCALE_ROUNDING_MODE: DequantScaleRoundingMode = DequantScaleRoundingMode.ROUND_UP):
    """
    Converts the src tensor to the output format specified by out_quant_type.
      axis: The axis along which the tensors are contiguous and quantization is applied.
      DEQUANT_SCALE_ROUNDING_MODE: 0 for ROUND_UP, 1 for ROUND_DOWN.

    Returns:
      out_quant_tensor: Quantized tensor in mx format.
         • For mxfp8, the output has the same shape as src_tensor.
         • For FP4, the size along the axis is halved, and the tensor is returned as a torch.uint8.
      scale: Scale tensor computed per microblock along the axis.

    If expected_scale is provided, apply the same preconditioning as float_to_flex
    before microscaled output quantization by dividing the source tensor by it.
    """
    # This should probably be packed into its own tiny class
    ndim = src_tensor.ndim
    assert -ndim <= axis < ndim, f"Invalid axis {axis=}"
    assert src_tensor.dtype in {torch.float32, torch.bfloat16,
                                torch.float16}, f"Invalid input tensor dtype {src_tensor.dtype}"

    axis = axis if axis >= 0 else axis + ndim
    is_fp4 = out_quant_type == torch.uint8
    is_fp8 = "float8" in str(out_quant_type)
    assert is_fp4 or is_fp8, f"Invalid input tensor dtype {out_quant_type}"
    assert scale_dtype == torch.uint8 or scale_dtype == torch.float8_e4m3fn, f"Invalid scale dtype {scale_dtype=}"
    if scale_dtype == torch.float8_e4m3fn:
        assert is_fp4, f"Direct float8 scales are only supported for FP4 values. Got {out_quant_type=}"
        assert DEQUANT_SCALE_ROUNDING_MODE == DequantScaleRoundingMode.ROUND_UP, \
            "Direct float8 scales only support ROUND_UP in downcast_to_mxfp_torch"

    device = src_tensor.device

    # For mxfp4 conversion, we assume the contiguous axis length is even.
    if is_fp4:
        axis_shape = src_tensor.size(axis)
        assert axis_shape % 2 == 0, "For mxfp4 conversion the contiguous axis length must be even."

    # Permute the tensor so that the contiguous axis becomes the last dimension.
    src = src_tensor.transpose(axis, src_tensor.ndim - 1).to(torch.float32)
    if expected_scale is not None:
        if expected_scale.numel() == 1:
            src = src / expected_scale.to(device=device, dtype=src.dtype)
        else:
            assert src.ndim == 3
            src = src / expected_scale.to(device=device, dtype=src.dtype)[:, None, None]
    axis_shape = src.shape[-1]

    # Pad the axis to be divisible by the microblock size, in case it is not.
    next_multiple = triton.cdiv(axis_shape, microblock_size) * microblock_size
    pad_amount = next_multiple - axis_shape
    padded_src = F.pad(src, (0, pad_amount))
    valid_mask = F.pad(torch.ones_like(src, dtype=torch.bool), (0, pad_amount))
    padded_axis_shape = padded_src.size(-1)

    # --- Compute per-group maximums for scale ---
    # Set padded entries to -1 so they don’t affect the max.
    abs_f = torch.abs(padded_src)
    abs_f = torch.where(valid_mask, abs_f, torch.tensor(-1.0, device=device, dtype=padded_src.dtype))
    # Reshape the last dimension into microblocks.
    new_shape = padded_src.shape[:-1] + (padded_axis_shape // microblock_size, microblock_size)
    abs_groups = abs_f.view(*new_shape)
    # Compute maximum along the microblock dimension.
    max_val, _ = abs_groups.max(dim=-1, keepdim=True)

    # Choose a max quantization value depending on type.
    max_quant_val = get_max_quant_val(out_quant_type)
    if DEQUANT_SCALE_ROUNDING_MODE == DequantScaleRoundingMode.ROUND_UP:
        dequant_scale = max_val / max_quant_val
    else:
        dequant_scale = max_val / (2 ** math.floor(math.log2(max_quant_val)))
    if scale_dtype == torch.uint8:
        # Convert to int to round the FP32 scale, prior to quantization!
        ds_int = dequant_scale.view(torch.int32)
        if DEQUANT_SCALE_ROUNDING_MODE == DequantScaleRoundingMode.ROUND_UP:
            ds_int_rounded = (ds_int + 0x007FFFFF) & 0x7F800000
        else:
            ds_int_rounded = ds_int & 0x7F800000
        dequant_scale_rounded = ds_int_rounded.view(torch.float32)
    else:
        # Direct fp8 scales keep the existing plain cast semantics here: the
        # stored scale is rounded by the fp8 conversion rather than forced
        # upward despite the ROUND_UP mode name.
        dequant_scale_rounded = dequant_scale.to(scale_dtype).to(torch.float32)

    # Compute the quantization scale.
    quant_scale = torch.where(dequant_scale_rounded == 0, torch.tensor(0.0, device=device), 1.0 / dequant_scale_rounded)

    # Quantize the tensor
    orig_padded_shape = padded_src.shape
    padded_src_groups = padded_src.view(*new_shape)
    quant_tensor = padded_src_groups * quant_scale
    # Reshape back to the original shape and trim padding
    quant_tensor = quant_tensor.view(orig_padded_shape)
    quant_tensor = quant_tensor[..., :axis_shape]

    # Finally, convert the quantized tensor to the target format
    if is_fp8:
        # Conversion must use satfinite PTX, so clamp before the conversion in torch to emulate this behavior
        quant_tensor = torch.clamp(quant_tensor, -max_quant_val, max_quant_val)
        out_weight = quant_tensor.to(out_quant_type)
    else:
        assert is_fp4, f"Invalid output quantization type {out_quant_type}"
        # For mxfp4, perform bit-level manipulation and pack two 4-bit values per uint8.
        # First, reinterpret the quantized tensor bits.
        q_int = quant_tensor.contiguous().view(torch.int32)
        # Extract sign, exponent, and mantissa.
        signs = q_int & 0x80000000
        exponents = right_shift_unsigned(q_int, 23) & 0xFF
        mantissas_orig = q_int & 0x7FFFFF

        E8_BIAS = 127
        E2_BIAS = 1
        # Adjust mantissas for subnormals.
        is_subnormal = exponents < E8_BIAS
        shift = E8_BIAS - exponents - 1
        mantissas_pre = (0x400000 | right_shift_unsigned(mantissas_orig, 1))
        bit0_dropped = (mantissas_orig & 0x1) != 0
        mask = (1 << shift.clamp(max=31)) - 1
        dropped_post = (mantissas_pre & mask) != 0
        sticky = is_subnormal & (bit0_dropped | dropped_post)
        mantissas = torch.where(is_subnormal, mantissas_pre >> shift, mantissas_orig)
        exponents = torch.maximum(exponents, torch.tensor(E8_BIAS - E2_BIAS, device=device)) - (E8_BIAS - E2_BIAS)
        # Round to nearest, ties to even (RTNE)
        m2bits = right_shift_unsigned(mantissas, 21) & 0x3
        lsb_keep = right_shift_unsigned(m2bits, 1) & 0x1
        guard = m2bits & 0x1
        sticky |= (mantissas & ((1 << 21) - 1)) != 0
        round_inc = guard & (sticky.to(torch.int32) | lsb_keep)
        e2m1_tmp = right_shift_unsigned(((exponents << 2) | m2bits) + round_inc, 1)
        e2m1_tmp = torch.minimum(e2m1_tmp, torch.tensor(0x7, device=device))
        e2m1_value = (right_shift_unsigned(signs, 28) | e2m1_tmp).to(torch.uint8)  # shape: (..., even_axis_shape)

        # Pack pairs of 4-bit values along the last dimension.
        e2m1_value = e2m1_value.view(*e2m1_value.shape[:-1], axis_shape // 2, 2)
        evens = e2m1_value[..., 0]
        odds = e2m1_value[..., 1]
        out_weight = evens | (odds << 4)  # shape: (..., axis_shape//2)

    # --- Process and output the scale ---
    if scale_dtype == torch.uint8:
        dq_scale = (ds_int_rounded.view(*dequant_scale.shape) >> 23).to(torch.uint8)
        dq_scale = dq_scale.squeeze(-1)
    else:
        dq_scale = dequant_scale_rounded.to(scale_dtype).squeeze(-1)
    out_weight = out_weight.transpose(axis, src_tensor.ndim - 1)
    dq_scale = dq_scale.transpose(axis, src_tensor.ndim - 1)
    return out_weight, dq_scale


def cvt_e2m1_to_fp32(input_tensor):
    assert input_tensor.dtype == torch.uint8

    input_tensor = input_tensor.to(torch.int32)
    evens = input_tensor & 0xF
    odds = (input_tensor >> 4) & 0xF

    vals = [0.0, 0.5, 1, 1.5, 2, 3, 4, 6]
    outputs = torch.tensor(vals, dtype=torch.float32, device=input_tensor.device)
    outputs = torch.cat([outputs, -outputs])

    even_floats = outputs[evens]
    odd_floats = outputs[odds]
    output_tensor = torch.stack([even_floats, odd_floats], dim=-1)
    output_tensor = output_tensor.view(*input_tensor.shape[:-1], input_tensor.shape[-1] * 2)
    return output_tensor


def upcast_from_mxfp_torch(tensor: torch.Tensor, scale: torch.Tensor, target_dtype: torch.dtype, axis: int):
    """
    Converts a microscaled tensor to the target format specified by target_dtype.
      axis: The axis along which dequantization is applied.

    Returns:
      out_weight: Tensor in the target format.
    """

    ndim = tensor.ndim
    assert -ndim <= axis < ndim, f"Invalid axis {axis=}"
    is_fp8 = tensor.dtype == torch.float8_e4m3fn or tensor.dtype == torch.float8_e5m2
    assert is_fp8 or tensor.dtype == torch.uint8, f"Invalid input quantization type {tensor.dtype}"
    assert scale.dtype == torch.uint8 or scale.dtype == torch.float8_e4m3fn, \
        f"Invalid scale dtype {scale.dtype}"

    # Permute the tensor and scale so that the quantization axis becomes the last dimension
    axis = axis if axis >= 0 else axis + ndim
    scale = scale.transpose(axis, scale.ndim - 1)
    tensor = tensor.transpose(axis, tensor.ndim - 1)

    if tensor.dtype == torch.uint8:
        fp32_tensor = cvt_e2m1_to_fp32(tensor)
    else:
        fp32_tensor = tensor.to(torch.float32)

    logical_quant_dim = tensor.shape[-1] * (2 if tensor.dtype == torch.uint8 else 1)
    if logical_quant_dim == 0:
        return fp32_tensor.to(target_dtype).transpose(axis, tensor.ndim - 1).contiguous()
    if scale.dtype == torch.uint8:
        dq_scale = (scale.to(torch.int32) << 23).view(torch.float32)
        scale_block_size = MXFP_BLOCK_SIZE.value
    else:
        dq_scale = scale.to(torch.float32)
        scale_block_size = logical_quant_dim // scale.shape[-1]
        assert scale_block_size in (NVFP_BLOCK_SIZE.value, MXFP_BLOCK_SIZE.value), f"Unsupported direct scale block size {scale_block_size}"
    axis_shape = fp32_tensor.size(-1)
    padded_axis_shape = triton.cdiv(logical_quant_dim, scale_block_size) * scale_block_size
    pad_size = padded_axis_shape - axis_shape
    padded_tensor = F.pad(fp32_tensor, (0, pad_size))

    new_axis_shape = padded_tensor.shape[-1]
    new_shape = padded_tensor.shape[:-1] + (new_axis_shape // scale_block_size, scale_block_size)
    padded_tensor = padded_tensor.view(*new_shape)
    dq_scale_padded = dq_scale.unsqueeze(-1)  # shape: [..., ceil(axis_shape/32), 1]
    out_padded = padded_tensor * dq_scale_padded
    # Need to clamp since due to rounding, we can have overflow that was within
    # the range before quantization.
    # e.g., 3.3895e+38 -> log2(3.3895e+38 / max_fp8e4m3=448) ~= 119.17 -> round
    # up to 120 + exp_bias=127 -> scale=247
    # 3.3895e+38 / 2**120 ~= 254.9976 -> round to 256 in fp8e4m3fn
    # Dequantization: 256 * 2**120 > 3.4e38 overflowing 3.38953139e38
    finfo = torch.finfo(target_dtype)
    out_padded = (padded_tensor * dq_scale_padded).clamp(finfo.min, finfo.max)
    if tensor.dtype == torch.float8_e5m2:
        # fp8e5m2 can have inf and we want to preserve so separately handle
        out_padded = out_padded.where(~padded_tensor.isinf(), padded_tensor.to(target_dtype))

    # Flatten back and remove the padded tail
    out_padded = out_padded.view(*fp32_tensor.shape[:-1], new_axis_shape)
    out_tensor = out_padded[..., :axis_shape]

    out_tensor = out_tensor.to(target_dtype).contiguous()
    out_tensor = out_tensor.transpose(axis, tensor.ndim - 1)

    return out_tensor


quantize_mxfp8_fn = _quantize_mxfp8_fn
quantize_mxfp4_fn = _quantize_mxfp4_fn
quantize_nvfp4_fn = _quantize_nvfp4_fn
