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


class MixedQuantizerFloat(nn.Module):
    r"""This is an mixedquantizer 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,
                 **kwargs):
        super(MixedQuantizerFloat, self).__init__()
        self.scale = torch.nn.Parameter(scale, requires_grad=False)
        self.zero_point = torch.nn.Parameter(zero_point, requires_grad=False)
        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

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

    def forward(self, x):
        if self.is_per_channel:
            if is_tracing_state():
                x = MixedQuantizePerchannelFloat.apply(
                    x, self.scale, self.zero_point, self.quant_min, self.quant_max, self.ch_axis)
            else:
                x = mixedquantize_function_per_channel_float(
                    x, self.scale, self.zero_point, self.quant_min, self.quant_max, self.ch_axis)
        else:
            if is_tracing_state():
                x = MixedQuantizePertensorFloat.apply(
                    x, self.scale, self.zero_point, self.quant_min, self.quant_max)
            else:
                x = mixedquantize_function_per_tensor_float(
                    x, self.scale, self.zero_point, self.quant_min, self.quant_max)
        return x



def mixedquantize_function_per_tensor_float(x, scale, zero_point, quant_min, quant_max):
    x = (x - zero_point) * scale
    x = x.round()
    x = torch.clamp(x, quant_min, quant_max)
    return x


def mixedquantize_function_per_channel_float(x, scale, 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)

    x = (x - zero_point) * scale
    x = x.round()
    x = torch.clamp(x, quant_min, quant_max)
    return x


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

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


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

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