# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai
# SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai
# SPDX-License-Identifier: Apache-2.0
# Contact: qubitium@modelcloud.ai, x.com/qubitium

import torch
from torch import nn

from ...adapter.adapter import Adapter, Lora
from ...models._const import DEVICE, PLATFORM
from ...nn_modules.qlinear import AWQuantLinear
from ...quantization import FORMAT, METHOD
from ...utils.awq import (
    awq_fast_gemm_forward_prefill,
    awq_fast_gemv_forward_decode,
    awq_runtime_available,
    awq_runtime_error,
)
from ...utils.backend import BACKEND
from ...utils.gemv import calculate_zeros_width


def pack_intweight(unpacked_qweight, interleave, kstride):
    # unpacked_qweight: [N, K]
    N = unpacked_qweight.shape[0]
    K = unpacked_qweight.shape[1]

    Packed_Kernel = unpacked_qweight.cpu().numpy().reshape(N, K // 32, 32)
    # np.arange(32).reshape(4, 4, 2).transpose(1, 0, 2) => [0, 1, 8, 9, 16, 17, 24, 25, ...]
    Packed_Kernel = Packed_Kernel.reshape(N, K // 32, 4, 4, 2).transpose(0, 1, 3, 2, 4)
    Packed_Kernel = Packed_Kernel.reshape(N, K // 32, 32)

    # reorder each 8 weights for fast dequantization
    # [0, 1, 2, 3, 4, 5, 6, 7] => [0, 2, 4, 6, 1, 3, 5, 7]
    Packed_Kernel = Packed_Kernel.reshape(N, K // 32, 4, 8)
    Packed_Kernel = Packed_Kernel.reshape(N, K // 32, 4, 4, 2).transpose(0, 1, 2, 4, 3)
    Packed_Kernel = Packed_Kernel.reshape(N, K)

    # interleaving every four rows
    Packed_Kernel = Packed_Kernel.reshape(
        N // interleave, interleave, K // kstride, kstride
    )
    # N // 4, K // 64, 4, 64
    Packed_Kernel = Packed_Kernel.transpose(0, 2, 1, 3)
    Packed_Kernel = Packed_Kernel.reshape(
        N // interleave, K // kstride, kstride, interleave
    )
    # Packing -> (N // 4, K // 64, 64)
    Packed_Kernel = (
            Packed_Kernel[..., 0]
            | (Packed_Kernel[..., 1] << 4)
            | (Packed_Kernel[..., 2] << 8)
            | (Packed_Kernel[..., 3] << 12)
    )
    # reshape to (N // 4, K), FP16 format
    Packed_Kernel = Packed_Kernel.reshape(N // interleave, K)
    qweight = (
        torch.tensor(Packed_Kernel.astype("int16"))
        .to(unpacked_qweight.device)
        .contiguous()
    )
    return qweight


class AwqGEMVFastLinear(AWQuantLinear):
    SUPPORTS_BACKENDS = [BACKEND.AWQ_GEMV_FAST]
    SUPPORTS_METHODS = [METHOD.AWQ]
    SUPPORTS_FORMATS = {FORMAT.GEMV_FAST: 30}
    SUPPORTS_BITS = [4]
    SUPPORTS_GROUP_SIZE = [-1, 16, 32, 64, 128]
    SUPPORTS_DESC_ACT = [True, False]
    SUPPORTS_SYM = [True, False]
    SUPPORTS_SHARDS = True
    SUPPORTS_TRAINING = True
    SUPPORTS_AUTO_PADDING = False
    SUPPORTS_IN_FEATURES_DIVISIBLE_BY = [1]
    SUPPORTS_OUT_FEATURES_DIVISIBLE_BY = [1]

    SUPPORTS_DEVICES = [DEVICE.CUDA, DEVICE.ROCM]
    SUPPORTS_PLATFORM = [PLATFORM.ALL]
    SUPPORTS_PACK_DTYPES = [torch.int16]
    SUPPORTS_ADAPTERS = [Lora]

    SUPPORTS_DTYPES = [torch.float16]

    # for transformers/optimum tests compat
    QUANT_TYPE = "awq_gemv_fast"

    zeros_name = "qzeros"

    def __init__(
            self,
            bits: int,
            group_size: int,
            sym: bool,
            desc_act: bool,
            in_features: int,
            out_features: int,
            bias: bool = False,
            pack_dtype: torch.dtype = torch.int16,
            adapter: Adapter = None,
            register_buffers: bool = False,
            **kwargs,
    ):
        backend = kwargs.pop("backend", BACKEND.AWQ_GEMV_FAST)
        super().__init__(
            bits=bits,
            group_size=group_size,
            sym=sym,
            desc_act=desc_act,
            in_features=in_features,
            out_features=out_features,
            bias=bias,
            pack_dtype=pack_dtype,
            backend=backend,
            adapter=adapter,
            register_buffers=False,
            **kwargs)

        self.split_k_iters = 8
        self.interleave = 4

        # The GEMV_FAST packing is a bit special; qweight and scales/zeros do not use the same pack_dtype.
        # qweight is packaged using torch.int16.
        # scales/zeros are packaged using torch.int32.
        int32_pack_factor = 32 // self.bits

        if register_buffers:
            self.register_buffer(
                "qweight",
                torch.zeros((out_features // self.interleave, in_features // self.pack_factor * self.interleave),
                            dtype=self.pack_dtype),
            )
            self.register_buffer(
                self.zeros_name,
                torch.zeros(
                    calculate_zeros_width(in_features, self.group_size) * int32_pack_factor,
                    out_features,
                    dtype=torch.float16,
                ),
            )
            self.register_buffer(
                "scales",
                torch.zeros(
                    calculate_zeros_width(in_features, self.group_size) * int32_pack_factor,
                    out_features,
                    dtype=torch.float16,
                ),
            )

            if bias:
                self.register_buffer("bias", torch.zeros(out_features, dtype=torch.float16))
            else:
                self.bias = None

    def forward(self, x: torch.Tensor):
        if self.input_rows(x) == 0:
            return self.empty_linear_output(x)
        if not awq_runtime_available():
            raise ModuleNotFoundError("AWQ torch.ops kernels are not properly installed. Error: " + awq_runtime_error())

        inputs = x
        inputs_dim = inputs.dim()

        batch_size = None
        n_tokens = None
        # Some linear layers in certain models (for example, opt-12m model.decoder.layers.0.fc1) receive inputs of shape `[B*S, H]`,
        # so it's necessary to check if the dimension is equal to 3.
        if inputs_dim == 3:
            batch_size, n_tokens, _ = inputs.shape

        input_dtype = inputs.dtype
        if input_dtype != torch.float16:
            inputs = inputs.to(dtype=torch.float16)
        if not inputs.is_contiguous():
            inputs = inputs.contiguous()

        self._ensure_runtime_buffers(device=inputs.device, dtype=inputs.dtype)

        zeros = self._runtime_zeros()
        if inputs_dim == 3 and batch_size < 8 and n_tokens == 1:
            out = awq_fast_gemv_forward_decode(
                inputs,
                self.qweight,
                self.scales,
                zeros,
                inputs.numel() // inputs.shape[-1],
                self.out_features,
                self.in_features,
                self.group_size,
            )
        else:
            out = awq_fast_gemm_forward_prefill(
                inputs, self.qweight, self.scales, zeros
            )

        if input_dtype != torch.float16:
            out = out.to(dtype=input_dtype)

        out = out + self.bias if self.bias is not None else out

        if self.adapter:
            out = self.adapter.apply(x=x, out=out)

        return out

    def _ensure_runtime_buffers(self, *, device: torch.device, dtype: torch.dtype):
        if self.qweight.device != device or not self.qweight.is_contiguous():
            self.qweight = self.qweight.to(device=device).contiguous()

        zeros = self._runtime_zeros()
        if zeros.device != device or zeros.dtype != dtype or not zeros.is_contiguous():
            zeros = zeros.to(device=device, dtype=dtype).contiguous()
            if self.zeros_name == "qzeros":
                self.qzeros = zeros
            elif self.zeros_name == "scaled_zeros":
                self.scaled_zeros = zeros
            else:
                raise ValueError(f"Unsupported zeros buffer: {self.zeros_name}")

        if self.scales.device != device or self.scales.dtype != dtype or not self.scales.is_contiguous():
            self.scales = self.scales.to(device=device, dtype=dtype).contiguous()

        if self.bias is not None and (
            self.bias.device != device or self.bias.dtype != dtype or not self.bias.is_contiguous()
        ):
            self.bias = self.bias.to(device=device, dtype=dtype).contiguous()

    def _runtime_zeros(self) -> torch.Tensor:
        if self.zeros_name == "qzeros":
            return self.qzeros
        if self.zeros_name == "scaled_zeros":
            return self.scaled_zeros
        raise ValueError(f"Unsupported zeros buffer: {self.zeros_name}")

    def pack(self, linear: nn.Module, scales: torch.Tensor, zeros: torch.Tensor, g_idx: torch.Tensor = None):
        # need scales and zeros info for real quantization
        assert scales is not None and zeros is not None
        scale_zeros = zeros * scales

        pack_num = 32 // self.bits
        qscales = torch.zeros(
            (
                scales.shape[0],
                calculate_zeros_width(linear.in_features, self.group_size) * pack_num,
            ),
            dtype=torch.float16,
            device=scales.device,
        )
        qscales[:, : scales.shape[1]] = scales
        self.register_buffer("scales", qscales.transpose(1, 0).contiguous())
        if linear.bias is not None:
            self.register_buffer("bias", linear.bias.clone().half())
        else:
            self.bias = None

        intweight = []
        for idx in range(self.in_features):
            intweight.append(
                torch.round(
                    (linear.weight.data[:, idx] + scale_zeros[:, idx // self.group_size])
                    / qscales[:, idx // self.group_size]
                ).to(torch.int)[:, None]
            )
        intweight = torch.cat(intweight, dim=1)
        intweight = intweight.to(dtype=torch.int32)
        self.register_buffer("qweight", pack_intweight(
            intweight.contiguous(), interleave=4, kstride=64
        ))

        zeros = zeros.to(dtype=torch.int32)
        qzeros = torch.zeros_like(qscales)

        qzeros[:, : scales.shape[1]] = -(
                qscales[:, : scales.shape[1]] * (zeros.to(torch.float32))
        ).to(torch.float16)
        self.register_buffer(self.zeros_name, qzeros.transpose(1, 0).contiguous())

    def extra_repr(self) -> str:
        return (
            "in_features={}, out_features={}, bias={}, bits={}, group_size={}".format(
                self.in_features,
                self.out_features,
                self.bias is not None,
                self.bits,
                self.group_size,
            )
        )


class LLMAwqLinear(AwqGEMVFastLinear):
    SUPPORTS_BACKENDS = [BACKEND.AWQ_GEMV_FAST]
    SUPPORTS_METHODS = [METHOD.AWQ]
    SUPPORTS_FORMATS = {FORMAT.LLM_AWQ: 100}
    SUPPORTS_BITS = [4]
    SUPPORTS_GROUP_SIZE = [-1, 16, 32, 64, 128]
    SUPPORTS_DESC_ACT = [True, False]
    SUPPORTS_SYM = [True, False]
    SUPPORTS_SHARDS = True
    SUPPORTS_TRAINING = True
    SUPPORTS_AUTO_PADDING = False
    SUPPORTS_IN_FEATURES_DIVISIBLE_BY = [1]
    SUPPORTS_OUT_FEATURES_DIVISIBLE_BY = [1]

    SUPPORTS_DEVICES = [DEVICE.CUDA, DEVICE.ROCM]
    SUPPORTS_PLATFORM = [PLATFORM.ALL]
    SUPPORTS_PACK_DTYPES = [torch.int16]
    SUPPORTS_ADAPTERS = [Lora]

    SUPPORTS_DTYPES = [torch.float16]

    # for transformers/optimum tests compat
    QUANT_TYPE = "llm-awq"

    zeros_name = "scaled_zeros"


__all__ = ["AwqGEMVFastLinear", "LLMAwqLinear"]