# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import torch
import math
from deepspeed.utils import logger
from deepspeed.ops.quantizer import ds_quantizer

TWO_D_PARAMS = 6


class Quantizer(object):

    def __init__(self,
                 q_groups=1,
                 q_mixed_fp16=False,
                 q_change_ratio=0.01,
                 q_type=0,
                 q_rounding=0,
                 q_verbose=False,
                 q_eigenvalue=False,
                 use_quantizer_kernel=False,
                 layer_num=0):

        self.q_groups = q_groups
        self.q_mixed_fp16 = q_mixed_fp16
        self.q_change_ratio = q_change_ratio
        self.q_type = q_type
        self.qsteps = 0
        self.quantize_real_ratio = 1.000
        self.q_verbose = q_verbose
        self.q_eigenvalue = q_eigenvalue
        self.use_quantizer_kernel = use_quantizer_kernel
        self.q_rounding = q_rounding
        self.layer_num = layer_num

    def any_precision_switch(self):
        # Temporary disabled functionality
        if self.layer_num == 0:
            return True
        result = False
        for index in range(self.layer_num):
            if self.q_start_bits[index] != self.q_target_bits:
                next_step = self.qsteps + (TWO_D_PARAMS * (self.layer_num if self.layer_num != 0 else 1))
                if next_step >= self.q_period[index]:
                    result = True
        return result

    def quantize(self, parameter_group, overflow, eigenvalue_enabled, block_eigenvalue={}):

        if overflow and not eigenvalue_enabled:
            return

        self.step()

        self.update_fp16_ratio()

        for i in range(len(parameter_group)):
            for p in parameter_group[i]:
                if len(p.size()) > 1 and hasattr(p, "start_bits") and p.start_bits:
                    param_id = id(p)
                    if block_eigenvalue is None:
                        eigenvalue, layer_id = None, 0
                    else:
                        eigenvalue, layer_id = block_eigenvalue[param_id] if param_id in block_eigenvalue else (None,
                                                                                                                0)
                    if eigenvalue is not None:
                        factor = 1 + math.floor(eigenvalue * 4)
                        p.data = self.compute_quantization(p.data, layer_id, factor)
                    else:
                        p.data = self.compute_quantization(p, layer_id)

    def step(self):
        self.qsteps += 1

    def quantize_highbit(self, inputs, num_bits):

        q_range = 2**num_bits
        input_flat = inputs.reshape(self.q_groups, -1)
        g_min = input_flat.amin(dim=-1, keepdim=True)
        g_max = input_flat.amax(dim=-1, keepdim=True)

        # Random number generator (Uniform)
        if self.q_rounding == 'nearest':
            p = 0.
        else:
            p = input_flat.new(input_flat.shape).uniform_(-0.5, 0.5)

        if self.q_type == 'symmetric':
            scale = 2 * torch.max(torch.abs(g_min), torch.abs(g_max)) / q_range
            zero_point = 0.
            input_flat = (input_flat / scale + p).round().clamp(-(q_range >> 1), (q_range >> 1) - 1) * scale
        elif self.q_type == 'asymmetric':
            scale = (g_max - g_min) / q_range
            zero_point = (g_min / scale).round() * scale
            input_flat = ((input_flat - zero_point) / scale + p).round().clamp(0, (q_range - 1)) * scale + zero_point
        output = input_flat.reshape(inputs.shape).contiguous()
        return output

    def quantize_tenary(self, inputs):
        input_flat = inputs.reshape(self.q_groups, -1)
        n = input_flat.shape[1]
        m = input_flat.norm(p=1, dim=1).div(n)
        thres = (0.7 * m).view(-1, 1)  #.expand_as(input_flat)
        pos = (input_flat > thres).type(inputs.type())
        neg = (input_flat < -thres).type(inputs.type())
        mask = (input_flat.abs() > thres).type(inputs.type())
        alpha = ((mask * input_flat).abs().sum(dim=1) / mask.sum(dim=1)).view(-1, 1)
        output = alpha * pos - alpha * neg
        output = output.reshape(inputs.shape).contiguous()
        return output

    def quantize_binary(self, inputs):
        input_flat = inputs.reshape(self.q_groups, -1)
        n = input_flat.shape[1]
        m = input_flat.norm(p=1, dim=1, keepdim=True).div(n)
        output = input_flat.sign().mul(m)
        output = output.reshape(inputs.shape).contiguous()
        return output

    def mixed_fp16_quantize(self, input, input_q, index):
        if self.q_mixed_fp16 and self.q_start_bits[index] >= (self.q_target_bits - 1):
            input_q = input * self.quantize_real_ratio + (1 - self.quantize_real_ratio) * input_q
            return input_q
        return input_q

    def compute_quantization(self, input, index=0, factor=1):
        # fixing the quantization bits based on the training steps
        # when reducing 1 bit at each period, we increase the period
        # to go slowly toward the target quantization bits
        # the period and starting bit can be configured

        if input.start_bits != input.target_bits:
            if self.qsteps >= input.q_period:
                self.quantize_real_ratio = 1.0
                input.q_period <<= 1
                input.q_period *= factor
                input.start_bits -= 1
                if self.q_verbose:
                    logger.info(
                        f'Quantization settings: current bit-precision = {input.start_bits}, step = {self.qsteps}, quantization period = {input.q_period}, index = {index}'
                    )
        assert (input.start_bits >= input.target_bits), \
            'Quantization bit is lower than target precision bits!'

        if self.use_quantizer_kernel:
            if input.start_bits <= 2:
                raise ValueError('Quantization bit is too low, please do it without quantization kernel!')
            input_q = ds_quantizer(input.data.clone(),
                                   self.q_groups,
                                   input.start_bits,
                                   asym=False if self.q_type == 'symmetric' else True,
                                   sr=False if self.q_rounding == 'nearest_neighbor' else True)
        else:
            if input.start_bits >= 3:
                input_flat = self.quantize_highbit(input.data, input.start_bits)
            elif input.start_bits == 2:
                assert self.q_type == 'symmetric', 'Quantization type is not symmetric!'
                assert self.q_rounding == 'nearest', 'Quantization rounding is not nearest_neighbor!'
                input_flat = self.quantize_tenary(input.data)
            elif input.start_bits == 1:
                assert self.q_type == 'symmetric', 'Quantization type is not symmetric!'
                assert self.q_rounding == 'nearest', 'Quantization rounding is not nearest_neighbor!'
                input_flat = self.quantize_binary(input.data)
        if self.use_quantizer_kernel:
            return self.mixed_fp16_quantize(input.data, input_q, index)
        else:
            if self.q_mixed_fp16 and input.start_bits >= input.target_bits - 1:
                input_flat = self.quantize_real_ratio * input.data + \
                              (1 - self.quantize_real_ratio) * input_flat
            return input_flat

    def update_fp16_ratio(self):
        if self.q_mixed_fp16:
            if self.quantize_real_ratio > 0:
                self.quantize_real_ratio -= self.q_change_ratio
            else:
                self.quantize_real_ratio = 0.000