# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from typing import TYPE_CHECKING, Literal

import paddle
from paddle import _C_ops
from paddle.base.data_feeder import check_dtype
from paddle.device import (
    is_compiled_with_cuda,
)
from paddle.device.cuda import get_device_capability
from paddle.framework import (
    LayerHelper,
    in_dynamic_or_pir_mode,
)

if TYPE_CHECKING:
    from typing_extensions import TypeAlias

    from paddle import Tensor
    from paddle._typing import DTypeLike

    _Algo: TypeAlias = Literal[
        'weight_only_int8', 'weight_only_int4', 'llm.int8'
    ]
    _GroupSize: TypeAlias = Literal[-1, 64, 128]


def _get_arch_info():
    # Get SMVersion from device.
    if is_compiled_with_cuda():
        cuda_version = paddle.version.cuda()
        if (
            cuda_version is not None and cuda_version != 'False'
        ) or paddle.is_compiled_with_rocm():
            major, minor = get_device_capability()
            arch = int(major * 10 + minor)
            return arch
        else:
            raise ValueError(
                "Paddle is not compiled with CUDA, we cannot get SMVersion from device, please try to compile Paddle with CUDA"
            )
    else:
        # Default arch value for type checking.
        return 0


def weight_quantize(
    x: Tensor,
    algo: _Algo = "weight_only_int8",
    arch: int | None = None,
    group_size: _GroupSize = -1,
) -> tuple[Tensor, Tensor]:
    """
    Quantization function for weight_only and llm.int8's weight.

    Args:
        x (Tensor): The input Tensor to be quantized, the data type is float16 or bfloat16.
        algo (str): The algo that is x will be apply, must be one of 'weight_only_int8',
            'weight_only_int4', 'llm.int8', 'w4a8' and 'w4afp8, default: 'weight_only_int8'.
        arch (int): The compute arch for target device. For example, A100 is 80, v100 is 70, if you do not assign arch, we will get arch from your device, default: None.
        group_size (int): The group size for weight quantization. -1 stands for default per-channel mode. Currently only support 64 or 128.

    Returns:
        out (Tensor): The Tensor which is the quantitative results, the data type is int8, the shape is transposition of x.
        scale (Tensor): The scale Tensor which is the scale of pre-channel, the data type is float32.
    Examples:
        .. code-block:: pycon

            >>> # doctest: +SKIP('No testing required')
            >>> import paddle
            >>> from paddle.nn.quant import weight_quantize

            >>> paddle.seed(2023)
            >>> x = paddle.rand(shape=[64, 32], dtype=paddle.float16)
            >>> out, scale = weight_quantize(x, algo='weight_only_int8')
            >>> print(out.shape)
            paddle.Size([32, 64])
            >>> print(scale.shape)
            paddle.Size([32])
    """
    if arch is None:
        arch = _get_arch_info()

    if is_compiled_with_cuda():
        assert (
            arch == 70
            or arch == 75
            or arch == 80
            or arch == 86
            or arch == 89
            or arch == 90
            or arch == 92
            or arch == 100
        ), (
            f"Currently weight_quantize only support SM70/75/80/86/89/90/92/100. but got {arch} "
        )

    assert group_size == -1 or group_size == 64 or group_size == 128, (
        f"Currently group_size only support -1/64/128. but got {group_size} "
    )
    if in_dynamic_or_pir_mode():
        return _C_ops.weight_quantize(x, algo, arch, group_size)
    else:
        type = "weight_quantize"
        helper = LayerHelper(type, **locals())
        out = helper.create_variable_for_type_inference('int8')
        scale = helper.create_variable_for_type_inference('float')

        helper.append_op(
            type=type,
            inputs={"x": x},
            outputs={'out': out, "scale": scale},
            attrs={"algo": algo, "arch": arch, "group_size": group_size},
        )
        return (out, scale)


def weight_dequantize(
    x: Tensor,
    scale: Tensor,
    algo: _Algo = "weight_only_int8",
    out_dtype: DTypeLike = "float16",
    group_size: _GroupSize = -1,
) -> Tensor:
    """
    Dequantization function for weight_only and llm.int8's weight.

    Args:
        x (Tensor): The input Tensor to be dequantized, the data type is int8.
        scale (Tensor): The scale Tensor which is the output of weight_quantize, the data type is float32.
        algo (str): The algo that is x will be apply, must be one of 'weight_only_int8',
            'weight_only_int4' and 'llm.int8', default: 'weight_only_int8'.
        out_dtype (str|np.dtype): [Deprecated][Not used] The output Tensor's data type, must be one of 'float16' and 'bfloat16', default: 'float16'.

    Returns:
        out (Tensor): The Tensor which is the dequantitative results, the data type is float16 or bfloat16, the shape is transposition of x.

    Examples:
        .. code-block:: pycon

            >>> # doctest: +SKIP('No testing required')
            >>> import paddle
            >>> from paddle.nn.quant import weight_quantize, weight_dequantize

            >>> paddle.seed(2023)
            >>> x = paddle.rand(shape=[64, 32], dtype=paddle.float16)
            >>> out, scale = weight_quantize(x, algo='weight_only_int8')
            >>> x_dequant = weight_dequantize(out, scale)
    """
    assert group_size == -1 or group_size == 64 or group_size == 128, (
        f"Currently group_size only support -1/64/128. but got {group_size} "
    )

    if in_dynamic_or_pir_mode():
        return _C_ops.weight_dequantize(x, scale, algo, group_size)
    else:
        type = "weight_dequantize"
        helper = LayerHelper(type, **locals())
        out_dtype = scale.dtype
        out = helper.create_variable_for_type_inference(out_dtype)

        helper.append_op(
            type=type,
            inputs={"x": x, "scale": scale},
            outputs={'out': out},
            attrs={
                "algo": algo,
                "group_size": group_size,
            },
        )
        return out


def weight_only_linear(
    x: Tensor,
    weight: Tensor,
    bias: Tensor | None = None,
    weight_scale: Tensor | None = None,
    weight_dtype: DTypeLike = "int8",
    arch: int | None = None,
    group_size: _GroupSize = -1,
) -> Tensor:
    """
    Applies matrix multiplication of two tensors and then bias addition if provided.
    This method requires CUDA version >= 11.2.

    Args:
        x (Tensor): The first input Tensor to be multiplied, the data type is float16 or bfloat16.
        weight (Tensor): The second input Tensor to be multiplied. Its rank must be 2.
        bias (Tensor|None): The input bias Tensor. If it is None, no bias addition would
            be performed. Otherwise, The bias is added to the matrix multiplication result.
        weight_scale (Tensor|None): The input scale Tensor Provided to weight for dequantization. Its rank must be 1.
        weight_dtype(str): The dtype of  weight Tensor, must be one of 'int8', 'int4', Defaulted to 'int8'.
        arch (int): The compute arch for target device. For example, A100 is 80, v100 is 70, if you do not assign arch, we will get arch from your device, default: None.
        group_size (int): The group size for weight quantization. -1 stands for default per-channel mode. Currently only support 64 or 128.
    Returns:
        Tensor: the output Tensor, the data type is the same as that of x.

    Examples:
        .. code-block:: pycon

            >>> # doctest: +SKIP('No testing required')
            >>> import paddle
            >>> from paddle.nn.quant import weight_only_linear

            >>> x = paddle.cast(paddle.randn([1, 2, 64]), dtype='float16')
            >>> weight = paddle.cast(paddle.randint(0, 127, [32, 64]), dtype='int8')
            >>> scale = paddle.randn([32], dtype='float32')
            >>> bias = paddle.cast(paddle.randn([32]), dtype='float16')
            >>> if paddle.device.cuda.get_device_capability()[0] >= 8:
            ...     out = weight_only_linear(
            ...         x,
            ...         weight,
            ...         bias=bias,
            ...         weight_scale=scale,
            ...         weight_dtype='int8',
            ...     )
            ...     print(out.shape)
            paddle.Size([1, 2, 32])
    """
    if arch is None:
        arch = _get_arch_info()

    if is_compiled_with_cuda():
        assert (
            arch == 70
            or arch == 75
            or arch == 80
            or arch == 86
            or arch == 89
            or arch == 90
            or arch == 92
            or arch == 100
        ), (
            f"Currently weight_quantize only support SM70/75/80/86/89/90/92/100. but got {arch} "
        )
    assert group_size == -1 or group_size == 64 or group_size == 128, (
        f"Currently weight_quantize only support group size of -1, 64 or 128. but got {group_size} "
    )

    if in_dynamic_or_pir_mode():
        out = _C_ops.weight_only_linear(
            x, weight, bias, weight_scale, weight_dtype, arch, group_size
        )
        return out
    else:
        check_dtype(
            weight_dtype, 'weight_dtype', ['int8', 'int4'], 'weight_only_linear'
        )
        type = "weight_only_linear"
        helper = LayerHelper(type, **locals())
        dtype = x.dtype

        inputs = {
            'x': [x],
            'weight': [weight],
            'weight_scale': [weight_scale],
        }
        if bias is not None:
            inputs["bias"] = [bias]
        attrs = {
            'weight_dtype': weight_dtype,
            'arch': arch,
            'group_size': group_size,
        }

        out = helper.create_variable_for_type_inference(dtype)

        helper.append_op(
            type=type,
            inputs=inputs,
            outputs={'out': out},
            attrs=attrs,
        )
        return out


def llm_int8_linear(
    x: Tensor,
    weight: Tensor,
    bias: Tensor | None = None,
    weight_scale: Tensor | None = None,
    threshold: float = 6.0,
) -> Tensor:
    """
    Applies matrix multiplication of two tensors and then bias addition if provided.
    This method requires CUDA version >= 11.2.

    Args:
        x (Tensor): the first input Tensor to be multiplied, the data type is float16 or bfloat16.
        weight (Tensor): the second input Tensor to be multiplied. Its rank must be 2.
        bias (Tensor|None): the input bias Tensor. If it is None, no bias addition would
            be performed. Otherwise, the bias is added to the matrix multiplication result.
        weight_scale (Tensor|None): the input scale Tensor Provided to weight for dequantization. Its rank must be 1.
        threshold(float): The min value of outlier in activation, outlier's channel will be apply multiply with x.dtype.

    Returns:
        Tensor: the output Tensor, the data type is the same as that of x.

    Examples:
        .. code-block:: pycon

            >>> # doctest: +SKIP('No testing required')
            >>> import paddle
            >>> from paddle.nn.quant import llm_int8_linear

            >>> x = paddle.cast(paddle.randn([1, 2, 64]), dtype='float16')
            >>> weight = paddle.cast(paddle.randint(0, 127, [32, 64]), dtype='int8')
            >>> scale = paddle.randn([32], dtype='float32')
            >>> bias = paddle.cast(paddle.randn([32]), dtype='float16')
            >>> if paddle.device.cuda.get_device_capability()[0] >= 8:
            ...     out = llm_int8_linear(x, weight, bias=bias, weight_scale=scale, threshold=6.0)
            ...     print(out.shape)
            paddle.Size([1, 2, 32])
    """
    if in_dynamic_or_pir_mode():
        out = _C_ops.llm_int8_linear(x, weight, bias, weight_scale, threshold)
        return out
    else:
        type = "llm_int8_linear"
        helper = LayerHelper(type, **locals())
        dtype = x.dtype

        inputs = {
            'x': [x],
            'weight': [weight],
            'weight_scale': [weight_scale],
        }
        if bias:
            inputs["bias"] = [bias]
        attrs = {'threshold': threshold}

        out = helper.create_variable_for_type_inference(dtype)

        helper.append_op(
            type=type,
            inputs=inputs,
            outputs={'out': out},
            attrs=attrs,
        )
        return out


def apply_per_channel_scale(x: Tensor, scales: Tensor) -> Tensor:
    """
    Apply pre-quant per channel scale on activations

    Args:
        x (Tensor): Input tensor representing the activations, the data type can be float16 or bfloat16.
        scales(Tensor): Per-channel scale factors for pre-quantization. Data type should be compatible with x.

    Returns:
        out (Tensor): The Tensor which is the pre-quant results, the data type is compatible with x.

    Examples:
        .. code-block:: pycon

            >>> # doctest: +SKIP('No testing required')
            >>> import paddle
            >>> from paddle.nn.quant import apply_per_channel_scale

            >>> paddle.seed(2023)
            >>> x = paddle.rand(shape=[64, 32], dtype=paddle.float16)
            >>> scales = paddle.rand(shape=[32], dtype=paddle.float16)
            >>> out = apply_per_channel_scale(x, scales)
    """

    if in_dynamic_or_pir_mode():
        return _C_ops.apply_per_channel_scale(x, scales)
    else:
        type = "apply_per_channel_scale"
        helper = LayerHelper(type, **locals())
        out = helper.create_variable_for_type_inference(x.dtype)

        helper.append_op(
            type=type,
            inputs={"x": [x], "scales": [scales]},
            outputs={"out": out},
        )
        return out