import torch
import torch.nn as nn
import numpy as np
from mqbench.utils.utils import is_tracing_state


class QuantizerInt(nn.Module):
    r"""This is an quantizer module for simulator fix-point-only inference.
    note that for activation quantization with per-channel format is not supported.
    if is_per_channel == True, per-channel quantization is used for weight quantization
    """
    def __init__(self, 
                 scale=1, 
                 zero_point=0, 
                 is_per_channel=False, 
                 is_symmetric_quant=True, 
                 ch_axis=1, 
                 quant_min=None, 
                 quant_max=None,
                 scale_bit=15,
                 dynamic_shift=True,
                 **kwargs):
        super(QuantizerInt, self).__init__()
        self.is_per_channel = is_per_channel
        self.ch_axis = ch_axis
        self.quant_min = quant_min
        self.quant_max = quant_max
        self.is_symmetric_quant = is_symmetric_quant
        self.scale_bit = scale_bit
        self.dynamic_shift = dynamic_shift
        int_scale, shift = self._get_fixpoint_scale_shift(scale)
        self.scale = torch.nn.Parameter(int_scale, requires_grad=False)
        self.zero_point = torch.nn.Parameter(zero_point.to(torch.float32).round().to(torch.int32), requires_grad=False)
        self.shift = torch.nn.Parameter(shift, requires_grad=False)
        self.dtype = kwargs.get('dtype', torch.float32)

    def _get_fixpoint_scale_shift(self, scale):
        if self.dynamic_shift:
            # offset = -torch.ceil(torch.log2(torch.max(scale.abs()))).long()
            offset = -torch.ceil(torch.log2(scale.abs())).long()
        else:
            offset = 0
        shift = self.scale_bit + offset
        if self.quant_min < 0:
            # int format
            quant_max_thresh = 2 ** (self.scale_bit - 1) - 1
            quant_min_thresh = -2 ** (self.scale_bit - 1)
            _scale = 2 ** (shift - 1)
            shift -= 1
        else:
            # Uint format
            quant_max_thresh = 2 ** (self.scale_bit) - 1
            quant_min_thresh = 0
            _scale = 2 ** shift

        scale = scale * _scale
        scale = torch.round(scale)
        scale = torch.clamp(scale, min=quant_min_thresh, max=quant_max_thresh).long()
        return scale, shift

    @torch.jit.export
    def extra_repr(self):
        return 'quant_min={}, quant_max={}, is_per_channel={}, is_symmetric_quant={}, ch_axis={}, ' \
               'scale_bit={}, is_dynamic_shift={}, scale={}, zero_point={}, shift={}'.format(
                   self.quant_min, self.quant_max, self.is_per_channel, self.is_symmetric_quant, self.ch_axis,
                   self.scale_bit,
                   self.dynamic_shift,
                   self.scale,
                   self.zero_point,
                   self.shift,)

    def forward(self, x):
        # x = x.to(self.dtype)
        if self.is_per_channel:
            if is_tracing_state():
                x = QuantizePerchannelInt.apply(
                    x, self.scale, self.shift, self.zero_point, self.quant_min, self.quant_max, self.ch_axis)
            else:
                x = quantize_function_per_channel_int(
                    x, self.scale, self.shift, self.zero_point, self.quant_min, self.quant_max, self.ch_axis)
        else:
            if is_tracing_state():
                x = QuantizePertensorInt.apply(
                    x, self.scale, self.shift, self.zero_point, self.quant_min, self.quant_max)
            else:
                x = quantize_function_per_tensor_int(
                    x, self.scale, self.shift, self.zero_point, self.quant_min, self.quant_max)
        
        return x



def quantize_function_per_tensor_int(x, scale, shift, zero_point, quant_min, quant_max):
    x = x / scale * 2**shift + zero_point
    x = (x + 0.5).floor()
    x = torch.clamp(x, quant_min, quant_max)
    return x


def quantize_function_per_channel_int(x, scale, shift, zero_point, quant_min, quant_max, ch_axis):
    new_shape = [1] * len(x.shape)
    new_shape[ch_axis] = x.shape[ch_axis]
    scale = scale.reshape(new_shape)
    zero_point = zero_point.reshape(new_shape)
    shift = shift.reshape(new_shape)

    x = x / scale * 2**shift + zero_point
    x = (x + 0.5).floor()
    x = torch.clamp(x, quant_min, quant_max)
    return x


class QuantizePertensorInt(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale, shift, zero_point, quant_min, quant_max):
        return quantize_function_per_tensor_int(x, scale, shift, zero_point, quant_min, quant_max)

    @staticmethod
    def symbolic(g, x, scale, shift, zero_point, quant_min, quant_max):
        return g.op("::QuantizePertensorInt", x, scale, shift, zero_point, quant_min_i=quant_min, quant_max_i=quant_max)


class QuantizePerchannelInt(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale, shift, zero_point, quant_min, quant_max, ch_axis):
        return quantize_function_per_channel_int(x, scale, shift, zero_point, quant_min, quant_max, ch_axis)

    @staticmethod
    def symbolic(g, x, scale, shift, zero_point, quant_min, quant_max, ch_axis):
        return g.op("::QuantizePerchannelInt", x, scale, shift, zero_point, quant_min_i=quant_min, quant_max_i=quant_max)