''' 量化模块
'''

from torch.nn import Linear, Embedding
from torch.nn.parameter import Parameter
from torch import nn
import torch.nn.functional as F
import bz2
import torch
import base64
import ctypes
from typing import List, Union, Dict
import re
from tqdm import tqdm
from functools import partial
import inspect
from bert4torch.snippets import is_package_available

try:
    from cpm_kernels.kernels.base import LazyKernelCModule, KernelFunction, round_up

    class Kernel:
        def __init__(self, code: bytes, function_names: List[str]):
            self.code = code
            self._function_names = function_names
            self._cmodule = LazyKernelCModule(self.code)

            for name in self._function_names:
                setattr(self, name, KernelFunction(self._cmodule, name))

    quantization_code = "$ "

    kernels = Kernel(
        bz2.decompress(base64.b64decode(quantization_code)),
        [
            "int4WeightCompression",
            "int4WeightExtractionFloat",
            "int4WeightExtractionHalf",
            "int8WeightExtractionFloat",
            "int8WeightExtractionHalf",
        ],
    )
except Exception as exception:
    kernels = None


class W8A16Linear(torch.autograd.Function):
    @staticmethod
    def forward(ctx, inp: torch.Tensor, quant_w: torch.Tensor, scale_w: torch.Tensor, weight_bit_width):
        ctx.inp_shape = inp.size()
        ctx.weight_bit_width = weight_bit_width
        out_features = quant_w.size(0)
        inp = inp.contiguous().view(-1, inp.size(-1))
        weight = extract_weight_to_half(quant_w, scale_w, weight_bit_width)
        ctx.weight_shape = weight.size()
        output = inp.mm(weight.t())
        ctx.save_for_backward(inp, quant_w, scale_w)
        return output.view(*(ctx.inp_shape[:-1] + (out_features,)))

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        inp, quant_w, scale_w = ctx.saved_tensors
        weight = extract_weight_to_half(quant_w, scale_w, ctx.weight_bit_width)
        grad_output = grad_output.contiguous().view(-1, weight.size(0))
        grad_input = grad_output.mm(weight)
        grad_weight = grad_output.t().mm(inp)
        return grad_input.view(ctx.inp_shape), grad_weight.view(ctx.weight_shape), None, None


def compress_int4_weight(weight: torch.Tensor):  # (n, m)
    with torch.cuda.device(weight.device):
        n, m = weight.size(0), weight.size(1)
        assert m % 2 == 0
        m = m // 2
        out = torch.empty(n, m, dtype=torch.int8, device="cuda")
        stream = torch.cuda.current_stream()

        gridDim = (n, 1, 1)
        blockDim = (min(round_up(m, 32), 1024), 1, 1)

        kernels.int4WeightCompression(
            gridDim,
            blockDim,
            0,
            stream,
            [ctypes.c_void_p(weight.data_ptr()), ctypes.c_void_p(out.data_ptr()), ctypes.c_int32(n), ctypes.c_int32(m)],
        )
        return out


def extract_weight_to_half(weight: torch.Tensor, scale_list: torch.Tensor, source_bit_width: int):
    if source_bit_width == 8:
        func = kernels.int8WeightExtractionHalf
    elif source_bit_width == 4:
        func = kernels.int4WeightExtractionHalf
    else:
        assert False, "Unsupported bit-width"

    with torch.cuda.device(weight.device):
        n, m = weight.size(0), weight.size(1)
        out = torch.empty(n, m * (8 // source_bit_width), dtype=torch.half, device="cuda")
        stream = torch.cuda.current_stream()

        gridDim = (n, 1, 1)
        blockDim = (min(round_up(m, 32), 1024), 1, 1)

        func(
            gridDim,
            blockDim,
            0,
            stream,
            [
                ctypes.c_void_p(weight.data_ptr()),
                ctypes.c_void_p(scale_list.data_ptr()),
                ctypes.c_void_p(out.data_ptr()),
                ctypes.c_int32(n),
                ctypes.c_int32(m),
            ],
        )
        return out


class CacheTensor():
    def __init__(self, *args, **kwargs):
        self.tensor = torch.empty(*args, **kwargs)

    def to(self, *args, **kwargs):
        self.tensor = self.tensor.to(*args, **kwargs)

    def data_ptr(self):
        return self.tensor.data_ptr()


class QuantizedLinear(Linear):
    def __init__(self, weight_bit_width: int, weight_tensor=None, bias_tensor=None, quantized_weight=None,
                 quantized_weight_scale=None, quantization_cache=None, empty_init=False, *args, **kwargs):
        super(QuantizedLinear, self).__init__(*args, **kwargs)
        self.weight_bit_width = weight_bit_width
        self.quantization_cache = quantization_cache

        if (quantized_weight is not None) and (quantized_weight_scale is not None):
            del self.weight
            self.weight = Parameter(quantized_weight.to(kwargs["device"]), requires_grad=False)
            self.weight_scale = Parameter(quantized_weight_scale.to(kwargs["device"]), requires_grad=False)
        else:
            shape = self.weight.shape
            del self.weight

            if weight_tensor is None or empty_init:
                self.weight = torch.empty(
                    shape[0], shape[1] * weight_bit_width // 8, dtype=torch.int8, device=kwargs["device"]
                )
                self.weight_scale = torch.empty(shape[0], dtype=kwargs["dtype"], device=kwargs["device"])
            else:
                self.weight_scale = (weight_tensor.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)).to(kwargs["dtype"])
                self.weight = torch.round(weight_tensor / self.weight_scale[:, None]).to(torch.int8)
                if weight_bit_width == 4:
                    self.weight = compress_int4_weight(self.weight)

            self.weight = Parameter(self.weight.to(kwargs["device"]), requires_grad=False)
            self.weight_scale = Parameter(self.weight_scale.to(kwargs["device"]), requires_grad=False)

        if bias_tensor is not None:
            self.bias = Parameter(bias_tensor.to(kwargs["device"]), requires_grad=False)
        else:
            self.bias = None

    def reset_parameters(self):
        """To accelerate initialization"""
        pass

    def forward(self, input):
        output = W8A16Linear.apply(input, self.weight, self.weight_scale, self.weight_bit_width)
        if self.bias is not None:
            output = output + self.bias
        return output

    def _apply(self, fn):
        self_obj = super()._apply(fn)
        if self.quantization_cache is not None:
            self.quantization_cache.to(self_obj.weight.device)
            self.quantization_cache.to(self_obj.weight_scale.dtype)
        return self_obj
    

class QuantizedEmbedding(Embedding):  # TODO: backward, check empty_init
    def __init__(self, weight_bit_width: int, weight_tensor=None, quantized_weight=None, quantized_weight_scale=None,
                 empty_init=False, *args, **kwargs):
        super(QuantizedEmbedding, self).__init__(*args, **kwargs)
        self.weight_bit_width = weight_bit_width

        if (quantized_weight is not None) and (quantized_weight_scale is not None):
            del self.weight
            self.weight = Parameter(quantized_weight.to(kwargs["device"]), requires_grad=False)
            self.weight_scale = Parameter(quantized_weight_scale.to(kwargs["device"]), requires_grad=False)
        else:
            shape = self.weight.shape
            del self.weight

            if weight_tensor is None or empty_init:
                self.weight = torch.empty(
                    shape[0], shape[1] * weight_bit_width // 8, dtype=torch.int8, device=kwargs["device"]
                )
                self.weight_scale = torch.empty(shape[0], dtype=kwargs["dtype"], device=kwargs["device"])
            else:
                self.weight_scale = (weight_tensor.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)).to(
                    kwargs["dtype"])
                self.weight = torch.round(weight_tensor / self.weight_scale[:, None]).to(torch.int8)
                if weight_bit_width == 4:
                    self.weight = compress_int4_weight(self.weight)

            self.weight = Parameter(self.weight.to(kwargs["device"]), requires_grad=False)
            self.weight_scale = Parameter(self.weight_scale.to(kwargs["device"]), requires_grad=False)

    def forward(self, input):
        original_weight = extract_weight_to_half(weight=self.weight, scale_list=self.weight_scale,
                                                    source_bit_width=self.weight_bit_width)
        output = F.embedding(
            input, original_weight, self.padding_idx, self.max_norm,
            self.norm_type, self.scale_grad_by_freq, self.sparse
        )
        return output


def quantize_cpm_kernels(model:nn.Module, quantization_bit:int=8, use_quantization_cache:bool=False, 
                         empty_init:bool=False, target_modules:Union[str, List]=None, **kwargs):
    """从chagglm-6b移植过来的的量化，方便以int8和int4进行推理
    源链接：https://huggingface.co/THUDM/chatglm-6b/blob/main/quantization.py
    
    Replace fp16 linear with quantized linear
    这里修改了hard code, 可以适配其他模型
    target_modules: str/list, 指定对某些层做量化
    """
    if not is_package_available('cpm_kernels'):
        raise ModuleNotFoundError('Module `cpm_kernels` not found, you may use `pip install cpm_kernels`')

    modules_trans = {}
    for name, module in model.named_modules():
        # target_modules=None, 表示对所有Linear层替换
        if (target_modules is None) and isinstance(module, Linear):
            modules_trans[name] = module
        elif (target_modules is not None) and isinstance(module, Linear):
            if isinstance(target_modules, str):
                target_module_found = re.fullmatch(target_modules, name)
            else:
                target_module_found = any(name.endswith(target_key) for target_key in target_modules)
            if target_module_found:
                modules_trans[name] = module
    
    # TODO: 暂时只支持cuda
    current_device = torch.cuda.current_device()
    dtype = torch.half

    QuantizedLinearWithPara = partial(
        QuantizedLinear,
        weight_bit_width=quantization_bit,
        bias=True,
        dtype=dtype,
        empty_init=empty_init
    )
        
    cache = dict()
    for name, module in tqdm(modules_trans.items(), desc='Quantize linear layers'):
        cache_name = re.sub(r'\.[0-9]+\.', '.', name)
        if use_quantization_cache and (cache_name not in cache):
            n, m = module.weight.size(0), module.weight.size(1)
            cache[cache_name] = CacheTensor(n, m, dtype=dtype, device=current_device, requires_grad=False)

        module_quant = QuantizedLinearWithPara(
                weight_tensor=module.weight.to(current_device),
                bias_tensor=module.bias,
                in_features=module.in_features,
                out_features=module.out_features,
                device=module.weight.device,
                quantization_cache=cache.get(cache_name)
            )
        del module
        # 赋值
        name_new = list(name)
        for iter_ in re.finditer(r'\.[0-9]+\.', name):
            iter_str = name[iter_.start():iter_.end()]
            name_new[iter_.start():iter_.end()] = [''] * (iter_.end()-iter_.start())
            name_new[iter_.start()] = '[' + iter_str[1:-1] + '].'
        exec('model.' + ''.join(name_new) + ' = module_quant')
    return model


def quantize_load_in_kbit(model:nn.Module, load_in_8bit:bool=False, load_in_4bit:bool=False, keep_in_fp32_modules:List=None, 
                          llm_int8_skip_modules:List=None, quantization_config:Dict=None, **kwargs):
    '''transformer的load_in_8bit, 源自transformer源代码'''
    # 兼容transformers新旧版本
    try:
        from transformers.integrations import replace_with_bnb_linear, set_module_quantized_tensor_to_device
    except:
        from transformers.utils.bitsandbytes import replace_with_bnb_linear, set_module_quantized_tensor_to_device
    from transformers.utils.quantization_config import BitsAndBytesConfig
    if quantization_config is None:
        quantization_config, kwargs = BitsAndBytesConfig.from_dict(
            config_dict={"load_in_8bit": load_in_8bit, "load_in_4bit": load_in_4bit}, return_unused_kwargs=True, **kwargs
        )
    elif quantization_config is not None:
        load_in_8bit = quantization_config.load_in_8bit
        load_in_4bit = quantization_config.load_in_4bit
        quantization_config_kwargs = {
            k: v for k, v in kwargs.items() if k in inspect.signature(BitsAndBytesConfig).parameters
        }

        if len(quantization_config_kwargs) > 0:
            raise ValueError(
                "You can't pass `load_in_8bit` or any other `BitsAndBytesConfig` argument as a kwarg when passing "
                "`quantization_config` argument at the same time."
            )

    load_in_8bit_skip_modules = quantization_config.llm_int8_skip_modules or []

    # We keep some modules such as the lm_head in their original dtype for numerical stability reasons
    modules_to_not_convert = load_in_8bit_skip_modules
    if not isinstance(modules_to_not_convert, list):
        modules_to_not_convert = [modules_to_not_convert]

    modules_to_not_convert.extend([] if keep_in_fp32_modules is None else keep_in_fp32_modules)
    modules_to_not_convert.extend([] if llm_int8_skip_modules is None else llm_int8_skip_modules)

    state_dict = model.state_dict()
    model = replace_with_bnb_linear(model, modules_to_not_convert=modules_to_not_convert, quantization_config=quantization_config)

    for key, param in model.named_parameters():
        if param.device == torch.device("meta"):
            set_module_quantized_tensor_to_device(model, key, 'cpu', value=state_dict[key])

    model.is_loaded_in_8bit = load_in_8bit
    model.is_loaded_in_4bit = load_in_4bit
    return model
