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

from fractions import Fraction
from typing import TYPE_CHECKING, Any

import torch

from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe import (
    RoutedExperts,
    UnquantizedFusedMoEMethod,
)
from vllm.model_executor.layers.linear import (
    LinearBase,
    UnquantizedLinearMethod,
)
from vllm.model_executor.layers.quantization import (
    QuantizationConfig,
    QuantizationMethods,
)
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

from .config_parser import INCConfigParser

if TYPE_CHECKING:
    from vllm.model_executor.models.utils import WeightsMapper

logger = init_logger(__name__)


class INCConfig(QuantizationConfig):
    """Config class for Intel Neural Compressor (INC).
    Repo: https://github.com/intel/neural-compressor
    """

    SUPPORTED_BITS = {2, 3, 4, 8}
    SUPPORTED_DTYPES = {"int", "mx_fp"}
    SUPPORTED_FORMATS = {
        "auto_round:auto_gptq",
        "auto_round:auto_awq",
        "auto_round:llm_compressor",
    }
    SUPPORTED_BACKENDS = {
        "auto",
        "gptq",
        "gptq:marlin",
        "awq",
        "awq:marlin",
        "marlin",
    }
    MXFP8_BITS = 8
    MXFP8_GROUP_SIZE = 32
    MXFP8_DATA_TYPE = "mx_fp"
    MXFP8_PACKING_FORMAT = "auto_round:llm_compressor"
    MXFP8_SUPPORTED_ACT_DTYPES = {"mx_fp", "mx_fp_rceil"}

    def __init__(
        self,
        weight_bits: int,
        group_size: int,
        sym: bool = True,
        packing_format: str = "auto_round:auto_gptq",
        block_name_to_quantize: str | list[str] | None = None,
        extra_config: dict[str, Any] | None = 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}."
            )
        # auto-round mxfp data_type is e.g. "mx_fp4" / "mx_fp4e2m1"; match the
        # "mx_fp" family by substring like auto_round.compressors.is_mx_fp.
        is_mxfp = "mx_fp" in data_type
        if data_type not in self.SUPPORTED_DTYPES and not is_mxfp:
            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)
        self.config_parser = INCConfigParser(self)

        self._validate_supported_quantization()

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

    @property
    def is_mxfp(self) -> bool:
        # auto-round mxfp data_type is the "mx_fp" family (e.g. "mx_fp",
        # "mx_fp4e2m1"); match by substring like auto_round.compressors.is_mx_fp.
        return "mx_fp" in self.data_type

    @property
    def is_mxfp8(self) -> bool:
        # MXFP4 and MXFP8 share data_type "mx_fp" and differ only by bit width.
        return self.is_mxfp and self.weight_bits == self.MXFP8_BITS

    def _validate_supported_quantization(self) -> None:
        if self.is_mxfp8:
            assert self.group_size == self.MXFP8_GROUP_SIZE, (
                "INC MXFP8 only supports group_size=32, "
                f"but found group_size={self.group_size}."
            )
            assert self.sym, "INC MXFP8 only supports symmetric weights."
            assert self.packing_format == self.MXFP8_PACKING_FORMAT, (
                "INC MXFP8 only supports "
                f"packing_format={self.MXFP8_PACKING_FORMAT!r}, "
                f"but found {self.packing_format!r}."
            )
            assert self.backend == "auto", (
                "INC MXFP8 only supports backend='auto', "
                f"but found backend={self.backend!r}."
            )
        elif self.packing_format == self.MXFP8_PACKING_FORMAT and not self.is_mxfp:
            raise ValueError(
                f"packing_format={self.MXFP8_PACKING_FORMAT!r} requires "
                f"an {self.MXFP8_DATA_TYPE!r} data_type."
            )

    def _validate_raw_config(self, config: dict[str, Any]) -> None:
        if not self.is_mxfp8:
            return

        expected_fields = {
            "act_bits": self.MXFP8_BITS,
            "act_data_type": self.MXFP8_DATA_TYPE,
            "act_group_size": self.MXFP8_GROUP_SIZE,
            "act_sym": True,
            "act_dynamic": True,
            "enable_quanted_input": False,
        }
        for field_name, expected_value in expected_fields.items():
            actual_value = self.get_from_keys_or(config, [field_name], expected_value)
            assert actual_value == expected_value, (
                "INC MXFP8 only supports "
                f"{field_name}={expected_value!r}, "
                f"but found {field_name}={actual_value!r}."
            )

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

    @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]) -> "INCConfig":
        quant_config = 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"),
        )
        quant_config._validate_raw_config(config)
        return quant_config

    def get_layer_config(self, layer, layer_name: str):
        return self.config_parser.get_layer_config(layer, layer_name)

    def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
        if self.block_name_to_quantize is not None:
            self.block_name_to_quantize = hf_to_vllm_mapper.apply_list(
                self.block_name_to_quantize
            )
        if self.extra_config is not None:
            self.extra_config = hf_to_vllm_mapper.apply_dict(self.extra_config)

    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
        from .schemes.factory import resolve_scheme

        # Match original: check model.-prefixed names for unquantized layers
        if prefix and self.extra_config:
            for layer_name in self.extra_config:
                if (
                    layer_name == prefix or layer_name == f"model.{prefix}"
                ) and self.extra_config[layer_name].get("bits", 16) >= 16:
                    if isinstance(layer, RoutedExperts):
                        return UnquantizedFusedMoEMethod(layer.moe_config)
                    return UnquantizedLinearMethod()

        layer_config = self.config_parser.resolve(layer, prefix)
        if not layer_config.quantized:
            if isinstance(layer, (LinearBase, ParallelLMHead)):
                return UnquantizedLinearMethod()
            if isinstance(layer, RoutedExperts):
                return UnquantizedFusedMoEMethod(layer.moe_config)
            return None

        logger.debug(
            "[%s] Type: %s, Bits: %s, Group Size: %s, Sym: %s",
            prefix,
            layer.__class__.__name__,
            layer_config.bits,
            layer_config.group_size,
            layer_config.sym,
        )

        scheme = resolve_scheme(layer_config)
        if isinstance(layer, (LinearBase, ParallelLMHead)):
            return scheme.get_linear_method(self, layer, prefix, layer_config)
        if isinstance(layer, RoutedExperts):
            return scheme.get_moe_method(self, layer, prefix, layer_config)
        return None

    @classmethod
    def override_quantization_method(
        cls, hf_quant_cfg, user_quant, hf_config=None
    ) -> "QuantizationMethods | None":
        """Override the `auto-round` method to `inc`."""
        is_auto_round_format = hf_quant_cfg.get("quant_method", None) == "auto-round"
        if is_auto_round_format:
            return cls.get_name()
        return None