"""
MXFP (Microscaling Floating Point) quantization implementation.
Supports various MX formats like fp4_e2m1, fp6, fp8, etc.
"""
import torch
from typing import Tuple, Optional
from utils import (
round_ste,
floor_ste,
ceil_ste,
reshape_pad_tensor_by_group_size,
revert_tensor_by_pad,
)

MXFP format cache: (ebits, mbits, emax, max_norm, min_norm)

MXFP_FORMAT_CACHE = {
"mx_int8": (0, 8, 0, 1.984375, 0),
"mx_int4": (0, 4, 0, 1.75, 0),
"mx_int2": (0, 2, 0, 1.0, 0),
"mx_fp8e5m2": (5, 4, 15, 57344.0, 6.103515625e-05),
"mx_fp8": (4, 5, 8, 448.0, 0.015625),
"mx_fp8e4m3": (4, 5, 8, 448.0, 0.015625),
"mx_fp6e3m2": (3, 4, 4, 28.0, 0.25),
"mx_fp6": (2, 5, 2, 7.5, 1.0),
"mx_fp6e2m3": (2, 5, 2, 7.5, 1.0),
"mx_fp4": (2, 3, 2, 6.0, 1.0),
"mx_fp4e2m1": (2, 3, 2, 6.0, 1.0),
"mx_float16": (5, 12, 15, 65504.0, 6.103515625e-05),
"mx_fp16": (5, 12, 15, 65504.0, 6.103515625e-05),
"mx_bfloat16": (8, 9, 127, 3.3895313892515355e38, 1.1754943508222875e-38),
"mx_bf16": (8, 9, 127, 3.3895313892515355e38, 1.1754943508222875e-38),
}

def quant_element(
tensor: torch.Tensor,
ebits: int,
mbits: int,
max_norm: float,
mantissa_rounding: str = "even"
) -> torch.Tensor:
"""
Quantize tensor elements to specified floating point format.

Args:
    tensor: Input tensor
    ebits: Number of exponent bits
    mbits: Number of mantissa bits (including implicit 1)
    max_norm: Maximum representable normal value
    mantissa_rounding: Rounding method ('even', 'nearest', 'floor', 'stochastic')

Returns:
    Quantized tensor
"""
if ebits != 0:
    private_exp = floor_ste(torch.log2(torch.abs(tensor) + (tensor == 0).type(tensor.dtype)))
    min_exp = -(2.0 ** float(ebits - 1)) + 2
    private_exp = private_exp.clip(min=min_exp)
else:
    private_exp = None

# Scale up so appropriate number of mbits are in the integer portion
tensor = (
    tensor * (2.0 ** float(mbits - 2))
    if private_exp is None
    else tensor / (2.0 ** private_exp.float()) * (2.0 ** float(mbits - 2))
)

# Apply rounding
if mantissa_rounding == "even":
    abs_tensor = torch.abs(tensor)
    mask_tensor = ((abs_tensor - 0.5) % 2 == torch.zeros_like(abs_tensor)).type(tensor.dtype)
    tensor = torch.sign(tensor) * (floor_ste(abs_tensor + 0.5) - mask_tensor)
elif mantissa_rounding == "nearest":
    tensor = torch.sign(tensor) * round_ste(torch.abs(tensor))
elif mantissa_rounding == "floor":
    tensor = torch.sign(tensor) * floor_ste(torch.abs(tensor))
elif mantissa_rounding == "stochastic":
    tensor = torch.sign(tensor) * floor_ste(torch.abs(tensor) + torch.rand_like(tensor, requires_grad=False))
else:
    raise ValueError(f"mantissa_rounding '{mantissa_rounding}' not supported. Use: even, nearest, floor, stochastic")

# Undo scaling
tensor = (
    tensor / (2.0 ** float(mbits - 2))
    if private_exp is None
    else tensor / (2.0 ** float(mbits - 2)) * (2.0 ** private_exp.float())
)

tensor = torch.clamp(tensor, min=-max_norm, max=max_norm)
return tensor

def quant_mxfp(
tensor: torch.Tensor,
bits: int = 4,
group_size: int = 32,
v: float = 0,
max_scale: float = 1.0,
mantissa_rounding: str = "nearest",
data_type: str = "mx_fp4e2m1",
scale_method: str = "default",
) -> Tuple[torch.Tensor, torch.Tensor, None]:
"""
MXFP quantization (quantize and dequantize).

Args:
    tensor: Input tensor to quantize
    bits: Number of bits for quantization (e.g., 4, 6, 8)
    group_size: Number of elements sharing scale
        - 0: per-tensor quantization
        - -1 or > D: per-token quantization
        - K (e.g., 32, 64, 128): group-wise quantization
    v: Rounding value perturbation (for tuning)
    max_scale: Maximum scale coefficient
    mantissa_rounding: Rounding method ('even', 'nearest', 'floor', 'stochastic')
    data_type: MX format type (e.g., 'mx_fp4e2m1', 'mx_fp8e4m3')
    scale_method: Scale computation method ('default', 'rceil')

Returns:
    qdq_result: Quantized and dequantized tensor
    shared_exp: Shared exponent (scale)
    None: Reserved for zero point (not used in MXFP)
"""
tensor, orig_shape, pad_len = reshape_pad_tensor_by_group_size(tensor, group_size)

# Resolve data type
data_type = data_type if data_type in MXFP_FORMAT_CACHE else "mx_fp" + str(bits)
if data_type not in MXFP_FORMAT_CACHE:
    raise ValueError(f"Unknown data_type: {data_type}. Available: {list(MXFP_FORMAT_CACHE.keys())}")

ebits, mbits, emax, max_norm, min_norm = MXFP_FORMAT_CACHE[data_type]
orig_dtype = tensor.dtype
tensor = tensor.to(torch.float32)

# Compute shared exponent (scale)
max_val, _ = torch.max(torch.abs(tensor), dim=-1, keepdim=True)
if isinstance(max_scale, torch.Tensor):
    max_val *= (max_scale.unsqueeze(dim=-1)).to(tensor.device)
else:
    max_val *= max_scale

if scale_method == "rceil":
    # Ceiling-based scale for tighter range
    shared_exp = torch.where(
        max_val == 0,
        torch.ones_like(max_val),
        ceil_ste(torch.log2(max_val / max_norm))
    )
else:
    # Default: floor-based scale
    shared_exp = torch.where(max_val == 0, torch.ones_like(max_val), torch.log2(max_val))
    shared_exp = floor_ste(shared_exp)
    shared_exp = shared_exp - emax

# Clamp shared exponent
scale_emax = 2.0 ** float(8 - 1) - 1
shared_exp = shared_exp.clamp(min=-scale_emax, max=scale_emax)

# Apply scale and quantize
scale = torch.pow(2.0, shared_exp.float())
tensor = tensor / scale + v
tensor = torch.clamp(tensor, min=-max_norm, max=max_norm)
tensor = quant_element(tensor, ebits, mbits, max_norm, mantissa_rounding)

# Dequantize
tensor = tensor * scale
tensor = revert_tensor_by_pad(tensor, orig_shape=orig_shape, pad_len=pad_len)

return tensor.to(orig_dtype), shared_exp.to(orig_dtype), None

def quant_mxfp8_static(
tensor: torch.Tensor,
data_type: str = "mx_fp8e4m3",
mantissa_rounding: str = "nearest",
) -> Tuple[torch.Tensor, torch.Tensor, None]:
"""
MXFP8 static per-tensor quantization with fixed minmax=1.

Unlike dynamic quantization which computes scale from tensor data,
this method fixes max_val=1.0 so the scale is data-independent (static).
The entire tensor shares a single scale (per-tensor granularity).

Args:
    tensor: Input tensor to quantize
    data_type: MX format type ('mx_fp8e4m3' or 'mx_fp8e5m2')
    mantissa_rounding: Rounding method ('even', 'nearest', 'floor', 'stochastic')

Returns:
    qdq_result: Quantized and dequantized tensor
    shared_exp: Shared exponent (scale), scalar for per-tensor
    None: Reserved for zero point
"""
if data_type not in MXFP_FORMAT_CACHE:
    raise ValueError(f"Unknown data_type: {data_type}. Available: {list(MXFP_FORMAT_CACHE.keys())}")

ebits, mbits, emax, max_norm, min_norm = MXFP_FORMAT_CACHE[data_type]
orig_dtype = tensor.dtype
orig_shape = tensor.shape
tensor = tensor.to(torch.float32).reshape(1, -1)
device = tensor.device
# Static quantization: fix max_val=1.0, do NOT compute from tensor data
max_val = torch.tensor([[1.0]], dtype=torch.float32)

# Compute shared exponent from fixed max_val
shared_exp = floor_ste(torch.log2(max_val))
shared_exp = shared_exp - emax

# Clamp shared exponent
scale_emax = 2.0 ** float(8 - 1) - 1
shared_exp = shared_exp.clamp(min=-scale_emax, max=scale_emax)

# Apply scale and quantize
scale = torch.pow(2.0, shared_exp.float()).to(device)
tensor = tensor / scale
tensor = torch.clamp(tensor, min=-max_norm, max=max_norm)
tensor = quant_element(tensor, ebits, mbits, max_norm, mantissa_rounding)

# Dequantize
tensor = tensor * scale
tensor = tensor.reshape(orig_shape)

return tensor.to(orig_dtype), shared_exp.to(orig_dtype), None

Convenience functions for common formats

def quant_mxfp4(
tensor: torch.Tensor,
group_size: int = -1,
mantissa_rounding: str = "nearest",
**kwargs
) -> Tuple[torch.Tensor, torch.Tensor, None]:
"""MXFP4 (E2M1) quantization."""
return quant_mxfp(
tensor, bits=4, group_size=group_size,
data_type="mx_fp4e2m1", mantissa_rounding=mantissa_rounding,
**kwargs
)

def quant_mxfp8(
tensor: torch.Tensor,
group_size: int = -1,
mantissa_rounding: str = "nearest",
data_type: str = "mx_fp8e4m3",
**kwargs
) -> Tuple[torch.Tensor, torch.Tensor, None]:
"""MXFP8 quantization."""
return quant_mxfp(
tensor, bits=8, group_size=group_size,
data_type=data_type, mantissa_rounding=mantissa_rounding,
**kwargs
)

def quant_mxfp_per_block(
tensor: torch.Tensor,
block_size: int = 16,
bits: int = 4,
v: float = 0,
max_scale: float = 1.0,
mantissa_rounding: str = "nearest",
data_type: str = "mx_fp4e2m1",
scale_method: str = "default",
) -> Tuple[torch.Tensor, torch.Tensor, None]:
"""
MXFP per-block quantization for BNSD format tensors.

Quantizes tensor with shared scale per block_size × block_size region
on the S and D dimensions.

Args:
    tensor: Input tensor in BNSD format (Batch, Num_heads, Seq_len, Head_dim)
    block_size: Size of square block for shared scale (e.g., 16, 32)
    bits: Number of bits for quantization (e.g., 4, 6, 8)
    v: Rounding value perturbation
    max_scale: Maximum scale coefficient
    mantissa_rounding: Rounding method ('even', 'nearest', 'floor', 'stochastic')
    data_type: MX format type (e.g., 'mx_fp4e2m1', 'mx_fp8e4m3')
    scale_method: Scale computation method ('default', 'rceil')

Returns:
    qdq_result: Quantized and dequantized tensor (BNSD format)
    shared_exp: Shared exponent (scale) per block, shape: (B, N, num_S_blocks, num_D_blocks)
    None: Reserved for zero point (not used in MXFP)

Example:
    >>> tensor = torch.randn(1, 32, 512, 128)  # BNSD
    >>> q_tensor, scale, _ = quant_mxfp_per_block(tensor, block_size=16, bits=4)
    >>> print(q_tensor.shape)  # (1, 32, 512, 128)
    >>> print(scale.shape)     # (1, 32, 32, 8)
"""
if data_type not in MXFP_FORMAT_CACHE:
    raise ValueError(f"Unknown data_type: {data_type}. Available: {list(MXFP_FORMAT_CACHE.keys())}")

ebits, mbits, emax, max_norm, min_norm = MXFP_FORMAT_CACHE[data_type]
orig_dtype = tensor.dtype
orig_shape = tensor.shape  # (B, N, S, D)

# Check tensor format
if len(orig_shape) != 4:
    raise ValueError(f"Expected 4D tensor (BNSD format), got shape: {orig_shape}")

B, N, S, D = orig_shape
tensor = tensor.to(torch.float32)

# Calculate padding for S and D dimensions
pad_S = (block_size - S % block_size) % block_size
pad_D = (block_size - D % block_size) % block_size

# Pad tensor: (B, N, S, D) -> (B, N, S+pad_S, D+pad_D)
if pad_S > 0 or pad_D > 0:
    tensor = torch.nn.functional.pad(tensor, (0, pad_D, 0, pad_S), value=0.0)

S_padded = S + pad_S
D_padded = D + pad_D
num_S_blocks = S_padded // block_size
num_D_blocks = D_padded // block_size

# Reshape to isolate blocks: (B, N, num_S_blocks, block_size, num_D_blocks, block_size)
tensor = tensor.reshape(B, N, num_S_blocks, block_size, num_D_blocks, block_size)

# Permute to group block elements together: (B, N, num_S_blocks, num_D_blocks, block_size, block_size)
tensor = tensor.permute(0, 1, 2, 4, 3, 5)

# Flatten block: (B, N, num_S_blocks, num_D_blocks, block_size * block_size)
tensor = tensor.reshape(B, N, num_S_blocks, num_D_blocks, block_size * block_size)

# Store this shape for later reconstruction
block_shape = (B, N, num_S_blocks, num_D_blocks, block_size * block_size)

# Compute shared exponent (scale) per block
max_val, _ = torch.max(torch.abs(tensor), dim=-1, keepdim=True)  # (B, N, num_S_blocks, num_D_blocks, 1)
if isinstance(max_scale, torch.Tensor):
    max_val *= max_scale.unsqueeze(dim=-1).to(tensor.device)
else:
    max_val *= max_scale
# scale_method = "rceil"
print(f"------------ {scale_method} -------------")
if scale_method == "rceil":
    shared_exp = torch.where(
        max_val == 0,
        torch.ones_like(max_val),
        ceil_ste(torch.log2(max_val / max_norm))
    )
else:
    shared_exp = torch.where(max_val == 0, torch.ones_like(max_val), torch.log2(max_val))
    shared_exp = floor_ste(shared_exp)
    shared_exp = shared_exp - emax

# Clamp shared exponent
scale_emax = 2.0 ** float(8 - 1) - 1
shared_exp = shared_exp.clamp(min=-scale_emax, max=scale_emax)

# Apply scale and quantize
scale = torch.pow(2.0, shared_exp.float())
tensor = tensor / scale + v
tensor = torch.clamp(tensor, min=-max_norm, max=max_norm)
# mantissa_rounding = 'even', 'nearest', 'floor', 'stochastic'
# mantissa_rounding = 'nearest'
print(f"============ {mantissa_rounding} =============")
tensor = quant_element(tensor, ebits, mbits, max_norm, mantissa_rounding)

# Dequantize
tensor = tensor * scale

# Reshape back: (B, N, num_S_blocks, num_D_blocks, block_size * block_size) -> (B, N, num_S_blocks, block_size, num_D_blocks, block_size)
tensor = tensor.reshape(B, N, num_S_blocks, num_D_blocks, block_size, block_size)

# Permute back: (B, N, num_S_blocks, num_D_blocks, block_size, block_size) -> (B, N, num_S_blocks, block_size, num_D_blocks, block_size)
tensor = tensor.permute(0, 1, 2, 4, 3, 5)

# Reshape to padded shape: (B, N, S_padded, D_padded)
tensor = tensor.reshape(B, N, S_padded, D_padded)

# Remove padding: (B, N, S_padded, D_padded) -> (B, N, S, D)
if pad_S > 0 or pad_D > 0:
    tensor = tensor[:, :, :S, :D]

# Squeeze shared_exp to remove last dimension: (B, N, num_S_blocks, num_D_blocks)
shared_exp = shared_exp.squeeze(-1)
return tensor.to(orig_dtype), shared_exp.to(orig_dtype), None