"""The per-tensor quantization config"""

import functools
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Type, Union  # noqa: UP035

import numpy as np
from tvm import DataType, DataTypeCode, IRModule, relax, runtime, te, tirx, topi
from tvm.relax.frontend import nn
from tvm.runtime import Tensor

from mlc_llm.loader import QuantizeMapping
from mlc_llm.nn import MixtralExperts
from mlc_llm.op import cutlass, extern
from mlc_llm.support import logging

from .utils import (
    apply_sharding,
    compile_quantize_func,
    convert_uint_packed_fp8_to_float,
    is_final_fc,
    is_moe_gate,
    pack_weight,
)

logger = logging.getLogger(__name__)


@dataclass
class PerTensorQuantize:
    """Configuration for per-tensor quantization"""

    name: str
    kind: str
    activation_dtype: Literal["float8_e4m3fn", "float8_e5m2"]
    weight_dtype: Literal["float8_e4m3fn", "float8_e5m2"]
    storage_dtype: Literal["uint32", "float8_e4m3fn", "float8_e5m2"]
    model_dtype: Literal["float16"]
    quantize_embedding: bool = True
    quantize_final_fc: bool = True
    quantize_linear: bool = True

    num_elem_per_storage: int = 0
    max_int_value: int = 0
    use_scale: bool = True
    # The calibration mode for quantization. If set to "inference", the model is built for
    # inference. This should be used after calibration is done.
    # If set to "max", the model is built for calibration that computes the scale using max value of
    # the activations.
    calibration_mode: Literal["inference", "max"] = "inference"
    tensor_parallel_shards: int = 1

    def __post_init__(self):
        assert self.kind == "per-tensor-quant"
        self.num_elem_per_storage = (
            DataType(self.storage_dtype).bits // DataType(self.weight_dtype).bits
        )
        self.max_int_value = int(tirx.max_value(self.weight_dtype).value)
        self._quantize_func_cache = {}

    def quantize_model(
        self,
        model: nn.Module,
        quant_map: QuantizeMapping,
        name_prefix: str,
        tensor_parallel_shards: int,
    ) -> nn.Module:
        """
        Quantize model with per-tensor quantization

        Parameters
        ----------
        model : nn.Module
            The non-quantized nn.Module.

        quant_map : QuantizeMapping
            The quantize mapping with name mapping and func mapping.

        name_prefix : str
            The name prefix for visited weight.

        tensor_parallel_shards : int
            The number of tensor parallel shards.

        Returns
        -------
        ret : nn.Module
            The quantized nn.Module.
        """

        self.tensor_parallel_shards = tensor_parallel_shards

        class _Mutator(nn.Mutator):
            def __init__(self, config: PerTensorQuantize, quant_map: QuantizeMapping) -> None:
                super().__init__()
                self.config = config
                self.quant_map = quant_map

            def visit_module(self, name: str, node: nn.Module) -> Any:
                """
                The visiting method for per-tensor quantization of nn.Module nodes.

                Parameters
                ----------
                name : str
                    The name of the current node.

                node : nn.Module
                    The current node of nn.Module to mutate.

                Returns
                ------
                ret_node: Any
                    The new node to replace current node.
                """
                weight_name = f"{name}.weight"
                param_names = (
                    [f"{name}.q_weight", f"{name}.q_scale"]
                    if self.config.use_scale
                    else [
                        f"{name}.q_weight",
                    ]
                )
                if (
                    isinstance(node, nn.Linear)
                    and self.config.quantize_linear
                    and (not is_final_fc(name) or self.config.quantize_final_fc)
                    and not is_moe_gate(name, node)
                ):
                    self.quant_map.param_map[weight_name] = param_names
                    self.quant_map.map_func[weight_name] = self.config.quantize_weight
                    op = PerTensorQuantizeLinear.from_linear(node, self.config, name)
                elif isinstance(node, nn.Embedding) and self.config.quantize_embedding:
                    self.quant_map.param_map[weight_name] = param_names
                    self.quant_map.map_func[weight_name] = self.config.quantize_weight
                    op = PerTensorQuantizeEmbedding.from_embedding(node, self.config)
                elif isinstance(node, MixtralExperts):
                    self.quant_map.param_map[weight_name] = param_names
                    self.quant_map.map_func[weight_name] = self.config.quantize_weight
                    op = PerTensorQuantizeMixtralExperts.from_mixtral_experts(
                        node, self.config, name
                    )
                else:
                    return self.visit(name, node)

                if hasattr(op, "q_calibration_scale") and op.q_calibration_scale:
                    # update quant_map for calibration scale
                    param_name = f"{name}.q_calibration_scale"
                    old_map_func = self.quant_map.map_func[weight_name]

                    def map_func(*args, **kwargs):
                        # placeholder for calibration scale, the actual value will be set after
                        # calibration.
                        scale = runtime.empty(
                            shape=op.q_calibration_scale.shape,
                            dtype=op.q_calibration_scale.dtype,
                        )
                        return [*old_map_func(*args, **kwargs), scale]

                    self.quant_map.param_map[weight_name].append(param_name)
                    self.quant_map.map_func[weight_name] = map_func
                return op

        model.to(dtype=self.model_dtype)
        mutator = _Mutator(self, quant_map)
        model = mutator.visit(name_prefix, model)
        return model

    def quantize_weight(self, weight) -> List[Tensor]:  # noqa: UP006
        """
        Quantize weight with per-tensor quantization.

        Parameters
        ----------
        weight : Tensor
            The weight to quantize.

        Returns
        -------
        ret : List[Tensor]
            The quantized weight and the scale if use_scale is True.
        """
        device = weight.device
        device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]

        def _create_quantize_func() -> IRModule:
            if DataType(self.weight_dtype).type_code in [
                DataTypeCode.Float8E4M3FN,
                DataTypeCode.Float8E5M2,
            ]:
                quantize_func = functools.partial(
                    self.quantize_float8,
                    quantize_dtype=self.weight_dtype,
                    storage_dtype=self.storage_dtype,
                )
            else:
                assert NotImplementedError()

            class Quantizer(nn.Module):
                """Quantizer module for per-tensor quantization."""

                def main(self, weight: nn.Tensor):
                    return quantize_func(weight)

            mod = Quantizer()
            mod, _ = mod.export_tvm(
                spec={"main": {"weight": nn.spec.Tensor(weight.shape, weight.dtype)}}
            )
            return mod

        key = f"({weight.shape}, {weight.dtype}, {device_type}"
        quantize_func = self._quantize_func_cache.get(key, None)
        if quantize_func is None:
            logger.info("Compiling quantize function for key: %s", key)
            quantize_func = compile_quantize_func(_create_quantize_func(), device)
            self._quantize_func_cache[key] = quantize_func
        return quantize_func(weight)

    def quantize_float8(
        self,
        tensor: nn.Tensor,
        quantize_dtype: str,
        storage_dtype: str,
    ) -> Union[Tuple[nn.Tensor], Tuple[nn.Tensor, nn.Tensor]]:  # noqa: UP006
        """Per-tensor quantization for weight tensor, defined in tensor expression."""

        if self.use_scale:
            # min_scaling_factor taken from TRT-LLM
            def _compute_scale(x: te.Tensor) -> te.Tensor:
                max_abs = topi.max(topi.abs(x))
                min_scaling_factor = tirx.const(
                    1.0 / (self.max_int_value * 512.0), self.model_dtype
                )
                scale = topi.maximum(
                    max_abs.astype(self.model_dtype) / self.max_int_value,
                    min_scaling_factor,
                ).astype("float32")
                scale = topi.expand_dims(scale, axis=0)
                return scale

            scale = nn.tensor_expr_op(_compute_scale, "compute_scale", args=[tensor])
        else:
            scale = None

        def _compute_quantized_tensor(weight: te.Tensor, scale: Optional[te.Tensor]) -> te.Tensor:
            elem_storage_dtype = (
                f"uint{DataType(quantize_dtype).bits}"
                if DataType(storage_dtype).type_code == DataTypeCode.UINT
                else quantize_dtype
            )
            scaled_tensor = te.compute(
                shape=weight.shape,
                fcompute=lambda *idx: tirx.Cast(
                    self.storage_dtype,
                    tirx.reinterpret(
                        elem_storage_dtype,
                        tirx.Cast(
                            quantize_dtype,
                            weight(*idx) / scale(0) if scale is not None else weight(*idx),
                        ),
                    ),
                ),
            )

            if quantize_dtype == self.storage_dtype:
                return scaled_tensor

            packed_weight = pack_weight(
                scaled_tensor,
                axis=-1,
                num_elem_per_storage=self.num_elem_per_storage,
                weight_dtype=self.weight_dtype,
                storage_dtype=self.storage_dtype,
            )

            return packed_weight

        quantized_tensor = nn.tensor_expr_op(
            _compute_quantized_tensor, "compute_quantized_tensor", args=[tensor, scale]
        )

        if self.use_scale:
            return quantized_tensor, scale
        return (quantized_tensor,)

    def _dequantize(
        self,
        q_weight: te.Tensor,
        scale: Optional[te.Tensor] = None,
        out_shape: Optional[Sequence[tirx.Expr]] = None,
    ) -> te.Tensor:
        if self.use_scale:
            assert scale is not None
        if DataType(self.weight_dtype).type_code in [
            DataTypeCode.Float8E4M3FN,
            DataTypeCode.Float8E5M2,
        ]:
            return self.dequantize_float8(q_weight, scale, self.weight_dtype, out_shape)
        raise NotImplementedError()

    def dequantize_float8(
        self,
        q_tensor: te.Tensor,
        scale: Optional[te.Tensor],
        quantize_dtype: str,
        out_shape: Optional[Sequence[tirx.Expr]] = None,
    ) -> te.Tensor:
        """Dequantize a fp8 tensor (input or weight) to higher-precision float."""
        if quantize_dtype != self.storage_dtype:
            dequantized_tensor = convert_uint_packed_fp8_to_float(
                q_tensor,
                self.num_elem_per_storage,
                self.storage_dtype,
                self.model_dtype,
                quantize_dtype,
                axis=-1,
                out_shape=out_shape,
            )
        else:
            dequantized_tensor = q_tensor.astype(self.model_dtype)
        if scale is not None:
            dequantized_tensor = dequantized_tensor * scale.astype(dequantized_tensor.dtype)
        return dequantized_tensor


class PerTensorQuantizeLinear(nn.Module):
    """An nn.Linear module with per-tensor quantization."""

    def __init__(
        self,
        in_features: int,
        out_features: Union[int, tirx.Var],
        config: PerTensorQuantize,
        name: str,
        bias: bool = True,
        out_dtype: Optional[str] = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.out_dtype = out_dtype or config.model_dtype
        self.config = config
        self.name = name
        self.q_weight = nn.Parameter(
            (out_features, tirx.ceildiv(in_features, config.num_elem_per_storage)),
            config.storage_dtype,
        )
        self.q_calibration_scale = None
        if config.use_scale:
            self.q_scale = nn.Parameter((1,), "float32")
            if config.calibration_mode == "inference":
                self.q_calibration_scale = nn.Parameter((1,), "float32")
        else:
            self.q_scale = None
        if bias:
            self.bias = nn.Parameter(
                (out_features,), config.model_dtype if out_dtype is None else out_dtype
            )
        else:
            self.bias = None

    @classmethod
    def from_linear(
        cls, src: nn.Linear, config: PerTensorQuantize, name: str
    ) -> "PerTensorQuantizeLinear":
        """
        Converts a non-quantized nn.Linear to a per-tensor quantized PerTensorQuantizeLinear

        Parameters
        ----------
        src : nn.Linear
            The non-quantized nn.Linear.

        config : PerTensorQuantize
            The per-tensor quantization config.

        name: str
            The name of the layer.

        Returns
        -------
        ret : PerTensorQuantizeLinear
            The per-tensor quantized PerTensorQuantizeLinear layer.
        """
        out_features, in_features = src.weight.shape
        quantized_linear = cls(
            in_features=in_features,
            out_features=out_features,
            config=config,
            name=name,
            bias=getattr(src, "bias", None) is not None,
            out_dtype=src.out_dtype,
        )
        if quantized_linear.bias is not None:
            quantized_linear.bias.attrs = src.bias.attrs
        if "shard_strategy" in src.weight.attrs:
            shard = src.weight.attrs["shard_strategy"]
            apply_sharding(shard, f"{shard.name}_q_weight", quantized_linear.q_weight)
            # scale doesn't need to be sharded since it's the same for all shards
        return quantized_linear

    def forward(self, x: nn.Tensor) -> nn.Tensor:
        """
        Forward method for per-tensor quantized linear layer.

        Parameters
        ----------
        x : nn.Tensor
            The input tensor.

        Returns
        -------
        ret : nn.Tensor
            The output tensor for the per-tensor quantized linear layer.
        """
        # Note: Use calibration scale when calibration is enabled
        if self.config.calibration_mode == "inference":
            if self.q_calibration_scale:
                x /= self.q_calibration_scale.astype(x.dtype)
            x_q = x.astype(self.config.activation_dtype)
            x_scale = self.q_calibration_scale
        elif self.config.calibration_mode == "max":
            _, x_scale = self.config.quantize_float8(
                x,
                quantize_dtype=self.config.activation_dtype,
                storage_dtype=self.config.storage_dtype,
            )
            if self.config.tensor_parallel_shards > 1:
                x_scale = nn.ccl_allreduce(x_scale, "max")
            x_scale = nn.extern(
                "mlc_llm.calibration_observer",
                [f"{self.name}.q_calibration_scale", "max", x_scale],
                out=nn.Tensor.placeholder(x_scale.shape, x_scale.dtype),
            )
            x_q = (x / x_scale.astype(x.dtype)).astype(self.config.activation_dtype)
            x = x_q.astype(self.config.model_dtype) * x_scale.astype(self.config.model_dtype)
        else:
            raise ValueError(f"Unknown calibration mode: {self.config.calibration_mode}")

        if (
            self.config.weight_dtype == self.config.storage_dtype
            and self.config.calibration_mode == "inference"
        ):
            if (
                extern.get_store().cutlass_gemm
                and functools.reduce(lambda x, y: x * y, x_q.shape[:-1]) != 1
            ):
                # Dispatch to cutlass kernel for gemm when cutlass is available.
                scale = (
                    x_scale * self.q_scale
                    if self.config.use_scale
                    else nn.wrap_nested(
                        relax.Constant(runtime.tensor(np.array([1.0]).astype("float32"))),
                        "scale",
                    )
                )
                return cutlass.fp8_gemm(
                    x_q,
                    self.q_weight,
                    scale,
                    self.config.weight_dtype,
                    self.config.model_dtype,
                )
            x = nn.op.matmul(x_q, nn.permute_dims(self.q_weight), out_dtype="float32")
            if self.config.use_scale:
                scale = x_scale * self.q_scale
                x = x * scale
            x = x.astype(self.out_dtype)
        else:
            w = nn.op.tensor_expr_op(
                lambda weight, scale: self.config._dequantize(
                    weight,
                    scale,
                    out_shape=[
                        (
                            tirx.IntImm("int64", self.out_features)
                            if isinstance(self.out_features, int)
                            else weight.shape[0]
                        ),
                        tirx.IntImm("int64", self.in_features),
                    ],
                ),
                "dequantize",
                args=[self.q_weight, self.q_scale],
            )
            x = nn.op.matmul(x, nn.permute_dims(w), out_dtype=self.out_dtype)
        if self.bias is not None:
            x = x + self.bias
        return x

    def to(self, dtype: Optional[str] = None) -> None:
        """
        Override to() such that we do not convert bias if there is an out_dtype.
        Otherwise, we might run into dtype mismatch when computing x + self.bias.
        """
        self.q_weight.to(dtype=dtype)
        if self.q_scale:
            self.q_scale.to(dtype=dtype)
        if self.bias is not None and self.out_dtype is None:
            self.bias.to(dtype=dtype)
        if dtype is not None and isinstance(getattr(self, "dtype", None), str):
            self.dtype = dtype


class PerTensorQuantizeEmbedding(nn.Module):
    """An nn.Embedding module with group quantization"""

    def __init__(self, num: Union[int, tirx.Var], dim: int, config: PerTensorQuantize):
        self.num = num
        self.dim = dim
        self.config = config
        self.q_weight = nn.Parameter(
            (num, tirx.ceildiv(dim, config.num_elem_per_storage)), config.storage_dtype
        )
        if self.config.use_scale:
            self.q_scale = nn.Parameter((1,), "float32")
        else:
            self.q_scale = None

    @staticmethod
    def from_embedding(
        embedding: nn.Embedding, config: PerTensorQuantize
    ) -> "PerTensorQuantizeEmbedding":
        """
        Converts a non-quantized nn.Embedding to a per-tensor quantized PerTensorQuantizeEmbedding

        Parameters
        ----------
        linear : nn.Embedding
            The non-quantized nn.Embedding.

        config : PerTensorQuantize
            The per-tensor quantization config.

        Returns
        -------
        ret : PerTensorQuantizeEmbedding
            The per-tensor quantized embedding layer.
        """
        num, dim = embedding.weight.shape
        return PerTensorQuantizeEmbedding(num, dim, config)

    def forward(self, x: nn.Tensor):
        """
        Forward method for per-tensor quantized embedding layer.

        Parameters
        ----------
        x : nn.Tensor
            The input tensor.

        Returns
        -------
        ret : nn.Tensor
            The output tensor for the embedding layer.
        """
        w = nn.op.tensor_expr_op(
            lambda weight, scale: self.config._dequantize(
                weight,
                scale,
                out_shape=[
                    (
                        tirx.IntImm("int64", self.num)
                        if isinstance(self.num, int)
                        else weight.shape[0]
                    ),
                    tirx.IntImm("int64", self.dim),
                ],
            ),
            "dequantize",
            args=[self.q_weight, self.q_scale],
        )
        if x.ndim == 1:
            return nn.op.take(w, x, axis=0)
        return nn.op.reshape(
            nn.op.take(w, nn.op.reshape(x, shape=[-1]), axis=0),
            shape=[*x.shape, self.dim],
        )

    def lm_head_forward(self, x: nn.Tensor):
        """The lm_head forwarding, which dequantizes the weight
        and multiplies it with the input tensor.

        Parameters
        ----------
        x : nn.Tensor
            The input tensor.

        Returns
        -------
        ret : nn.Tensor
            The output tensor for the lm_head layer.
        """
        w = nn.op.tensor_expr_op(
            lambda weight, scale: self.config._dequantize(
                weight,
                scale,
                out_shape=[
                    (
                        tirx.IntImm("int64", self.num)
                        if isinstance(self.num, int)
                        else weight.shape[0]
                    ),
                    tirx.IntImm("int64", self.dim),
                ],
            ),
            "dequantize",
            args=[self.q_weight, self.q_scale],
        )
        w = nn.op.permute_dims(w)
        return nn.op.matmul(x, w, out_dtype="float32")


class PerTensorQuantizeMixtralExperts(nn.Module):
    """An MixtralExperts module with group quantization"""

    _IMPL: ClassVar[Dict[str, Type["PerTensorQuantizeMixtralExperts"]]] = {}  # noqa: UP006

    def __init__(
        self,
        num_local_experts,
        in_features,
        out_features,
        config: PerTensorQuantize,
        name: str,
    ):
        self.num_local_experts = num_local_experts
        self.in_features = in_features
        self.out_features = out_features
        self.config = config
        self.name = name
        self.q_weight = nn.Parameter(
            (
                num_local_experts,
                out_features,
                tirx.ceildiv(in_features, config.num_elem_per_storage),
            ),
            config.storage_dtype,
        )
        self.q_calibration_scale = None
        if config.use_scale:
            self.q_scale = nn.Parameter((1,), "float32")
            if config.calibration_mode == "inference":
                self.q_calibration_scale = nn.Parameter((1,), "float32")
        else:
            self.q_scale = None

    @staticmethod
    def from_mixtral_experts(
        src: "MixtralExperts",
        config: PerTensorQuantize,
        name: str,
    ) -> "PerTensorQuantizeMixtralExperts":
        """
        Converts a non-quantized MixtralExperts to a per-tensor quantized
        PerTensorQuantizeMixtralExperts

        Parameters
        ----------
        src : MixtralExperts
            The non-quantized MixtralExperts

        config : PerTensorQuantize
            The per-tensor quantization config

        name: str
            The name of the layer.

        Returns
        -------
        ret : PerTensorQuantizeMixtralExperts
            The per-tensor quantized MixtralExperts layer
        """
        if DataType(config.weight_dtype).type_code in [
            DataTypeCode.Float8E4M3FN,
            DataTypeCode.Float8E5M2,
        ]:
            return PerTensorQuantizeMixtralExperts._IMPL["fp8"].from_mixtral_experts(
                src, config, name
            )
        raise NotImplementedError()

    def forward(self, x: nn.Tensor, indptr: nn.Tensor) -> nn.Tensor:
        """Forward method for per-tensor quantized mistral experts.

        Parameters
        ----------
        x : nn.Tensor
            The input tensor.

        indptr: nn.Tensor
            The indptr tensor

        Returns
        -------
        ret : nn.Tensor
            The output tensor for the per-tensor quantized mistral experts layer.
        """
        raise NotImplementedError()