# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from fractions import Fraction
from typing import Any, Optional, Union

import torch

from vllm.logger import init_logger
from vllm.model_executor.layers.linear import (LinearBase,
                                               UnquantizedLinearMethod)
from vllm.model_executor.layers.quantization import QuantizationMethods
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types

logger = init_logger(__name__)


class AutoRoundConfig(QuantizationConfig):
    """Config class for AutoRound.
    Reference: https://arxiv.org/pdf/2309.05516
    """

    SUPPORTED_BITS = {2, 3, 4, 8}
    SUPPORTED_DTYPES = {"int"}
    SUPPORTED_FORMATS = {"auto_round:auto_gptq", "auto_round:auto_awq"}
    SUPPORTED_BACKENDS = {
        "auto", "gptq", "gptq:marlin", "awq", "awq:marlin", "marlin", "ipex"
    }

    def __init__(
        self,
        weight_bits: int,
        group_size: int,
        sym: bool = True,
        packing_format: str = "auto_round:auto_gptq",
        block_name_to_quantize: Optional[Union[str, list[str]]] = None,
        extra_config: Optional[dict[str, Any]] = None,
        data_type: str = "int",
        backend: str = "auto",
    ) -> None:
        super().__init__()
        if weight_bits not in self.SUPPORTED_BITS:
            raise ValueError(f"Unsupported weight_bits: {weight_bits}, "
                             f"currently only support  {self.SUPPORTED_BITS}")
        if data_type not in self.SUPPORTED_DTYPES:
            raise ValueError(
                f"Unsupported data_type: {data_type},"
                f" currently only support  {self.SUPPORTED_DTYPES}")
        if packing_format not in self.SUPPORTED_FORMATS:
            raise ValueError(
                f"Unsupported packing_format: {packing_format}, "
                f"currently only support  {self.SUPPORTED_FORMATS}")
        if backend not in self.SUPPORTED_BACKENDS:
            raise ValueError(
                f"Unsupported backend: {backend},  "
                f"currently only support  {self.SUPPORTED_BACKENDS}")

        self.weight_bits = weight_bits
        self.group_size = group_size
        self.sym = sym
        self.packing_format = packing_format
        self.block_name_to_quantize = (block_name_to_quantize.split(",") if
                                       isinstance(block_name_to_quantize, str)
                                       else block_name_to_quantize)
        self.extra_config = extra_config
        self.data_type = data_type
        self.backend = backend
        self.pack_factor = Fraction(32, weight_bits)

    def __repr__(self) -> str:
        return (f"AutoRoundConfig(weight_bits={self.weight_bits}, "
                f"group_size={self.group_size}, sym={self.sym})")

    @classmethod
    def get_name(cls) -> QuantizationMethods:
        return "auto-round"

    @classmethod
    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
        return [torch.half, torch.bfloat16]

    @classmethod
    def get_min_capability(cls) -> int:
        return 60

    @classmethod
    def get_config_filenames(cls) -> list[str]:
        return ["quantization_config.json"]

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> "AutoRoundConfig":
        return cls(
            weight_bits=cls.get_from_keys(config, ["bits"]),
            group_size=cls.get_from_keys(config, ["group_size"]),
            sym=cls.get_from_keys(config, ["sym"]),
            packing_format=cls.get_from_keys_or(config, ["packing_format"],
                                                "auto_round:auto_gptq"),
            block_name_to_quantize=cls.get_from_keys_or(
                config, ["block_name_to_quantize", "to_quant_block_names"],
                None),
            extra_config=cls.get_from_keys_or(config, ["extra_config"], None),
            data_type=cls.get_from_keys_or(config, ["data_type"], "int"),
            backend=cls.get_from_keys_or(config, ["backend", "vllm_backend"],
                                         "auto"),
        )

    def get_layer_config(self, layer, layer_name: str):
        # Priority: extra_config > block_name_to_quantize > type fallback
        if self.extra_config and layer_name in self.extra_config:
            cfg = self.extra_config[layer_name]
            return cfg.get("bits", self.weight_bits), cfg.get(
                "group_size", self.group_size), cfg.get("sym", self.sym)

        quantized = True
        if self.block_name_to_quantize:
            quantized = any(
                layer_name.startswith(name)
                for name in self.block_name_to_quantize)
        elif isinstance(layer, ParallelLMHead):
            quantized = False

        return (self.weight_bits, self.group_size,
                self.sym) if quantized else (16, -1, True)

    def check_quantized(self, weight_bits: int) -> bool:
        return weight_bits < 16

    def apply_awq_quant_layer(self, layer, prefix: str, backend: str = "auto"):
        from vllm.model_executor.layers.fused_moe import FusedMoE
        from vllm.model_executor.layers.quantization.utils.marlin_utils import (
            check_marlin_supported, check_moe_marlin_supports_layer)

        weight_bits, group_size, sym = self.get_layer_config(layer, prefix)
        if not self.check_quantized(weight_bits):
            if isinstance(layer, (LinearBase, ParallelLMHead)):
                return UnquantizedLinearMethod()
            else:
                return None

        logger.debug("[%s] Type: %s, Bits: %s, Group Size: %s, Sym: %s",
                     prefix, layer.__class__.__name__, weight_bits, group_size,
                     sym)
        if backend == "auto" or "marlin" in backend:
            AWQ_TYPE_MAP = {
                4: scalar_types.uint4,
                8: scalar_types.uint8,
            }
            use_marlin = (weight_bits
                          in AWQ_TYPE_MAP) and check_marlin_supported(
                              AWQ_TYPE_MAP[weight_bits], group_size, not sym)

            if isinstance(layer, FusedMoE):
                use_marlin = use_marlin and check_moe_marlin_supports_layer(
                    layer, group_size)

        else:
            use_marlin = False
        if use_marlin:
            from vllm.model_executor.layers.quantization.awq_marlin import (
                AWQMarlinConfig, AWQMarlinLinearMethod, AWQMoEMethod)
            quant_args_marlin = AWQMarlinConfig(weight_bits=weight_bits,
                                                group_size=group_size,
                                                zero_point=not sym,
                                                lm_head_quantized=False,
                                                full_config={},
                                                modules_to_not_convert=[])
        else:
            from vllm.model_executor.layers.quantization.awq import (
                AWQConfig, AWQLinearMethod)
            quant_args = AWQConfig(
                weight_bits=weight_bits,
                group_size=group_size,
                zero_point=not sym,
            )

        if isinstance(layer, FusedMoE):
            if use_marlin:
                return AWQMoEMethod(quant_args_marlin)
            from vllm.model_executor.layers.quantization.moe_wna16 import (
                MoeWNA16Config)
            config = {
                "quant_method": "awq",
                "bits": weight_bits,
                "group_size": group_size,
                "zero_point": not sym,
                "lm_head": False,
            }
            return MoeWNA16Config.from_config(config).get_quant_method(
                layer, prefix)

        if isinstance(layer, (LinearBase, ParallelLMHead)):
            if use_marlin:
                return AWQMarlinLinearMethod(quant_args_marlin)
            else:
                return AWQLinearMethod(quant_args)
        return None

    def apply_gptq_quant_layer(self,
                               layer,
                               prefix: str,
                               backend: str = "auto"):
        from vllm.model_executor.layers.fused_moe import FusedMoE
        from vllm.model_executor.layers.quantization.utils.marlin_utils import (
            check_marlin_supported, check_moe_marlin_supports_layer)
        weight_bits, group_size, sym = self.get_layer_config(layer, prefix)
        if not self.check_quantized(weight_bits):
            if isinstance(layer, (LinearBase, ParallelLMHead)):
                return UnquantizedLinearMethod()
            else:
                return None

        logger.debug("[%s] Type: %s, Bits: %s, Group Size: %s, Sym: %s",
                     prefix, layer.__class__.__name__, weight_bits, group_size,
                     sym)
        if backend == "auto" or "marlin" in backend:
            GPTQ_TYPE_MAP = {
                (4, True): scalar_types.uint4b8,
                (8, True): scalar_types.uint8b128,
            }
            use_marlin = ((weight_bits, sym) in GPTQ_TYPE_MAP
                          and check_marlin_supported(
                              GPTQ_TYPE_MAP[(weight_bits, sym)],
                              group_size,
                              has_zp=not sym))
            if isinstance(layer, FusedMoE):
                use_marlin = use_marlin and check_moe_marlin_supports_layer(
                    layer, group_size)
        else:
            use_marlin = False
        if use_marlin:
            from vllm.model_executor.layers.quantization.gptq_marlin import (
                GPTQMarlinConfig, GPTQMarlinLinearMethod, GPTQMarlinMoEMethod)
            quant_args_marlin = GPTQMarlinConfig(weight_bits=weight_bits,
                                                 group_size=group_size,
                                                 is_sym=sym,
                                                 lm_head_quantized=False,
                                                 desc_act=False,
                                                 dynamic={},
                                                 full_config={})
        else:
            from vllm.model_executor.layers.quantization.gptq import (
                GPTQConfig, GPTQLinearMethod)
            quant_args = GPTQConfig(weight_bits=weight_bits,
                                    group_size=group_size,
                                    lm_head_quantized=False,
                                    desc_act=False,
                                    dynamic={})

        if isinstance(layer, FusedMoE):
            if use_marlin:
                from vllm.model_executor.layers.quantization.moe_wna16 import (
                    MoeWNA16Config)
                config = {
                    "quant_method": "gptq",
                    "bits": weight_bits,
                    "group_size": group_size,
                    "sym": sym,
                    "lm_head": False,
                }
                return MoeWNA16Config.from_config(config).get_quant_method(
                    layer, prefix)
            return GPTQMarlinMoEMethod(quant_args_marlin)

        if isinstance(layer, (LinearBase, ParallelLMHead)):
            if use_marlin:
                return GPTQMarlinLinearMethod(quant_args_marlin)
            else:
                return GPTQLinearMethod(quant_args)

        return None

    def apply_ipex_quant_layer(self, layer, prefix: str):
        weight_bits, group_size, sym = self.get_layer_config(layer, prefix)
        if not self.check_quantized(weight_bits):
            if isinstance(layer, (LinearBase, ParallelLMHead)):
                return UnquantizedLinearMethod()
            else:
                return None
        from vllm.model_executor.layers.quantization.ipex_quant import (
            IPEXAWQLinearMethod, IPEXConfig, IPEXGPTQLinearMethod)
        if isinstance(layer, (LinearBase, ParallelLMHead)):
            if "awq" in self.packing_format:
                config = IPEXConfig(method="awq",
                                    weight_bits=weight_bits,
                                    group_size=group_size)
                return IPEXAWQLinearMethod(config)
            elif "gptq" in self.packing_format:
                config = IPEXConfig(method="gptq",
                                    weight_bits=weight_bits,
                                    group_size=group_size)
                return IPEXGPTQLinearMethod(config)
            else:
                raise ValueError(
                    f"ipex backend only supports awq "
                    f"and gtpq format,but got {self.packing_format}")
        else:
            return None

    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
        if (current_platform.is_cpu() or current_platform.is_xpu()
                or self.backend == "ipex"):
            return self.apply_ipex_quant_layer(layer, prefix)
        if "gptq" in self.packing_format or "gptq" in self.backend:
            return self.apply_gptq_quant_layer(layer, prefix)
        if "awq" in self.packing_format or "awq" in self.backend:
            return self.apply_awq_quant_layer(layer, prefix)
