from typing import Optional, Union

import numpy as np
import tensorrt as trt
import torch
from tensorrt import ITensor as TRTTensor
from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
from torch.fx.node import Target
from torch_tensorrt._utils import is_tensorrt_version_supported
from torch_tensorrt.dynamo._SourceIR import SourceIR
from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
from torch_tensorrt.dynamo.conversion.converter_utils import (
    get_trt_tensor,
    set_layer_name,
)

if is_tensorrt_version_supported("10.8.0"):

    def quantize(
        ctx: ConversionContext,
        target: Target,
        source_ir: Optional[SourceIR],
        name: str,
        input_tensor: TRTTensor,
        block_size: int,
        amax: Union[np.ndarray, torch.Tensor],
        num_bits: int,
        exponent_bits: int,
        scale_num_bits: int,
        scale_exponent_bits: int,
    ) -> TRTTensor:
        """
        Adds quantize and dequantize ops (QDQ) which quantize to FP4 based
        on the output_type set and dequantizes them back.
        """
        if len(input_tensor.shape) not in (2, 3):
            raise ValueError(
                f"dynamic_block_quantize converter received an input of {input_tensor.shape} shape. Supported shapes: 2D or 3D"
            )
        with unset_fake_temporarily():
            axis = -1
            global_scale = _calculate_global_scale(ctx, name, amax)
            if ".weight_quantizer" in name:
                output = _static_double_quantize(
                    ctx,
                    target,
                    source_ir,
                    name,
                    input_tensor,
                    global_scale,
                    axis,
                )
            elif ".input_quantizer" in name:
                output = _dynamic_double_quantize(
                    ctx,
                    target,
                    source_ir,
                    name,
                    input_tensor,
                    global_scale,
                    axis,
                )
            else:
                raise ValueError(
                    f"quantizer received an input of {name}. Supported values: weight_quantizer | input_quantizer"
                )
            return output

    def _dynamic_double_quantize(
        ctx: ConversionContext,
        target: Target,
        source_ir: Optional[SourceIR],
        name: str,
        input_tensor: TRTTensor,
        global_scale: torch.Tensor,
        axis: int = -1,
        block_size: int = 16,
        output_type: trt.DataType = trt.DataType.FP4,
        scale_type: trt.DataType = trt.DataType.FP8,
    ) -> TRTTensor:
        """
        quantize input tensor to fp4
        Parameters:
            ctx: ConversionContext,
            target: Target,
            source_ir: Optional[SourceIR]
            name: str
            input_tensor : TRTTensor (On GPU)
                The input TRTTensor.
            global_scale : Tensor (On GPU)
                The global per-tensor scaling factor. It should contain only 1 element.
            axis : int
                The axis to quantize. Default is -1 (the last axis).
            block_size : int
                The block size for quantization. Default is 16.
            output_type : trt.DataType
                The data type for quantized data. Default is FP4.
            scale_type : trt.DataType
                The data type for block scale. Default is FP8.

        """
        global_scale = get_trt_tensor(ctx, global_scale, name + "_global_scale")

        if input_tensor.dtype not in [
            trt.DataType.HALF,
            trt.DataType.FLOAT,
            trt.DataType.BF16,
        ]:
            raise ValueError(
                f"Currently supported input tensor type is float16 | float32 | bfloat16, got Unsupported dtype: {input_tensor.dtype}"
            )
        # dynamic quantize input tensor to fp4
        dynamic_quantize_layer = ctx.net.add_dynamic_quantize(
            input_tensor,
            axis,
            block_size,
            output_type,
            scale_type,
        )
        dynamic_quantize_layer.set_input(1, global_scale)
        set_layer_name(
            dynamic_quantize_layer, target, name + "_dynamic_quantize", source_ir
        )
        quantized_data_in_fp4 = dynamic_quantize_layer.get_output(0)
        quantized_scale_in_fp8 = dynamic_quantize_layer.get_output(1)

        return _double_dequantize(
            ctx,
            target,
            source_ir,
            name,
            quantized_data_in_fp4,
            quantized_scale_in_fp8,
            global_scale,
            axis,
            input_tensor.dtype,
        )

    def _double_dequantize(
        ctx: ConversionContext,
        target: Target,
        source_ir: Optional[SourceIR],
        name: str,
        quantized_data_in_fp4: TRTTensor,
        quantized_scale_in_fp8: TRTTensor,
        global_scale: torch.Tensor,
        axis: int = -1,
        output_type: trt.DataType = trt.DataType.FLOAT,
    ) -> TRTTensor:
        """
        double dequantize will first dequantize scale from fp8 to orignal dtype(default is float32)
        and then dequantize data from fp4 to orignal dtype(default is float32)
        Parameters:
            ctx: ConversionContext,
            target: Target,
            source_ir: Optional[SourceIR]
            name: str
            quantized_data_in_fp4: TRTTensor
            quantized_scale_in_fp8: TRTTensor
            global_scale: torch.Tensor
            axis: int
            output_type: trt.DataType
        """
        # dequantize scale from fp8 to orignal dtype(default is float32)
        dequantize_scale_layer = ctx.net.add_dequantize(
            quantized_scale_in_fp8, global_scale, output_type
        )
        dequantize_scale_layer.axis = axis
        dequantize_scale_layer.to_type = output_type
        set_layer_name(
            dequantize_scale_layer, target, name + "_dequantize_scale", source_ir
        )
        dequantized_scale = dequantize_scale_layer.get_output(0)

        # dequantize quantized_data_in_fp4 from fp4 to orignal dtype(default is float32)
        dequantize_data_layer = ctx.net.add_dequantize(
            quantized_data_in_fp4, dequantized_scale, output_type
        )
        dequantize_data_layer.axis = axis
        dequantize_data_layer.to_type = output_type
        set_layer_name(
            dequantize_data_layer, target, name + "_dequantize_data", source_ir
        )
        dequantized_data = dequantize_data_layer.get_output(0)
        return dequantized_data

    def _static_double_quantize(
        ctx: ConversionContext,
        target: Target,
        source_ir: Optional[SourceIR],
        name: str,
        weights_tensor: torch.Tensor,
        global_scale: torch.Tensor,
        axis: int,
    ) -> TRTTensor:
        """
        Parameters:
            ctx: ConversionContext,
            target: Target,
            source_ir: Optional[SourceIR],
            name: str,
            weights_tensor : Tensor (On GPU)
                The input tensor for weights.
            global_scale : Tensor (On GPU)
                The global per-tensor scaling factor. It should contain only 1 element.
            axis: int
                The axis to quantize. Default is -1 (the last axis).
        Returns:
            quantized data tensor in fp4
        """

        import modelopt.torch.quantization.qtensor.nvfp4_tensor as nvfp4_tensor

        if weights_tensor.dtype == torch.float16:
            original_dtype = trt.DataType.HALF
        elif weights_tensor.dtype == torch.float32:
            original_dtype = trt.DataType.FLOAT
        elif weights_tensor.dtype == torch.bfloat16:
            original_dtype = trt.DataType.BF16
        else:
            raise ValueError(
                f"Currently supported weights tensor type is float16 | float32 | bfloat16, got Unsupported dtype: {weights_tensor.dtype}"
            )
        block_scale_fp8 = nvfp4_tensor.NVFP4QTensor.get_weights_scaling_factor(
            weights_tensor,
            16,
            global_scale,
        )[0]
        weights_tensor_fp4 = nvfp4_tensor.NVFP4QTensor.quantize(
            weights_tensor,
            16,
            block_scale_fp8,
            global_scale,
        )[0]._quantized_data

        block_scale_fp8 = get_trt_tensor(
            ctx,
            block_scale_fp8,
            name + "_block_scale_fp8",
            target_quantized_type=trt.DataType.FP8,
        )
        global_scale = get_trt_tensor(ctx, global_scale, name + "_global_scale")
        weights_tensor_fp4 = get_trt_tensor(
            ctx,
            weights_tensor_fp4,
            name + "_weights_fp4",
            target_quantized_type=trt.DataType.FP4,
        )

        dequantized_data = _double_dequantize(
            ctx,
            target,
            source_ir,
            name,
            weights_tensor_fp4,
            block_scale_fp8,
            global_scale,
            axis,
            original_dtype,
        )
        return dequantized_data

    def _calculate_global_scale(
        ctx: ConversionContext,
        name: str,
        amax: torch.Tensor,
    ) -> torch.Tensor:
        # calculate global scale (the global per-tensor scaling factor, should only contain 1 element)
        assert len(amax.shape) == 0, "amax should be a scalar"
        global_scale = amax / 6 / 448
        global_scale.masked_fill_(global_scale == 0, 1.0)
        return global_scale