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


class AddQuantizerInt(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, 
                 scale1=1, 
                 scale2=1, 
                 bias=0,
                 is_per_channel=False, 
                 is_symmetric_quant=True, 
                 ch_axis=1, 
                 quant_min=None, 
                 quant_max=None,
                 scale_bit=15,
                 dynamic_shift=True,
                 shift_reduce=4,
                 half_range=True,
                 shift_bit=6,
                 **kwargs):
        super(AddQuantizerInt, 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.shift_bit = shift_bit
        self.dynamic_shift = dynamic_shift
        self.half_range = half_range
        self.shift_reduce = torch.nn.Parameter(torch.tensor([shift_reduce], device=bias.device, dtype=torch.float32), requires_grad=False)
        self.per_channel_shift = kwargs.get('per_channel_shift', False)
        int_scale1, shift1 = self._get_fixpoint_scale_shift(scale1)
        int_scale2, shift2 = self._get_fixpoint_scale_shift(scale2)
        self.scale1 = torch.nn.Parameter(int_scale1.to(torch.float32), requires_grad=False)
        self.shift1 = torch.nn.Parameter(shift1.to(torch.float32), requires_grad=False)
        self.scale2 = torch.nn.Parameter(int_scale2.to(torch.float32), requires_grad=False)
        self.shift2 = torch.nn.Parameter(shift2.to(torch.float32), requires_grad=False)
        bias = torch.floor(bias * 2**self.shift_reduce + 0.5)
        self.bias = torch.nn.Parameter(bias, requires_grad=False)

    def _get_fixpoint_scale_shift(self, scale):
        if self.dynamic_shift:
            if self.per_channel_shift:
                offset = -torch.ceil(torch.log2(scale.abs())).long()
            else:
                offset = -torch.ceil(torch.log2(torch.max(scale.abs()).reshape([1]))).long()
        else:
            offset = torch.tensor(0)

        if self.half_range:
            shift = self.scale_bit + offset
            quant_max_thresh = 2 ** (self.scale_bit) - 1
            quant_min_thresh = 0
        else:
            shift = self.scale_bit + offset - 1
            quant_max_thresh = 2 ** (self.scale_bit-1) - 1
            quant_min_thresh = -2 ** (self.scale_bit-1)
        _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={}, scale1={}, shift1={}, scale2={}, shift2={}, bias={}, shift_reduce={}'.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.scale1,
                   self.shift1,
                   self.scale2,
                   self.shift2,
                   self.bias,
                   self.shift_reduce
                   )

    def forward(self, x1, x2):
        if self.is_per_channel:
            if is_tracing_state():
                x = AddQuantizePerchannelInt.apply(
                    x1, x2, self.scale1, self.shift1, self.scale2, self.shift2, self.bias, self.quant_min, self.quant_max, self.ch_axis, self.shift_reduce)
            else:
                x = addquantize_function_per_channel_int(
                    x1, x2, self.scale1, self.shift1, self.scale2, self.shift2, self.bias, self.quant_min, self.quant_max, self.ch_axis, self.shift_reduce)
        else:
            if is_tracing_state():
                x = AddQuantizePertensorInt.apply(
                    x1, x2, self.scale1, self.shift1, self.scale2, self.shift2, self.bias, self.quant_min, self.quant_max, self.shift_reduce)
            else:
                x = addquantize_function_per_tensor_int(
                    x1, x2, self.scale1, self.shift1, self.scale2, self.shift2, self.bias, self.quant_min, self.quant_max, self.shift_reduce)
        
        return x

def addquantize_function_per_tensor_int(x1, x2, scale1, shift1, scale2, shift2, bias, quant_min, quant_max, shift_reduce):
    x1 = (x1 * scale1 / 2**(shift1-shift_reduce) + 0.5).floor()
    x2 = (x2 * scale2 / 2**(shift2-shift_reduce) + 0.5).floor()
    x = x1 + x2 + bias
    x = x / 2**shift_reduce
    x = (x + 0.5).floor()
    x = torch.clamp(x, quant_min, quant_max)
    return x


def addquantize_function_per_channel_int(x1, x2, scale1, shift1, scale2, shift2, bias, quant_min, quant_max, ch_axis, shift_reduce):
    new_shape = [1] * len(x1.shape)
    new_shape[ch_axis] = -1
    scale1 = scale1.reshape(new_shape)
    shift1 = shift1.reshape(new_shape)
    scale2 = scale2.reshape(new_shape)
    shift2 = shift2.reshape(new_shape)
    bias = bias.reshape(new_shape)

    x1 = (x1 * scale1 / 2**(shift1-shift_reduce) + 0.5).floor()
    x2 = (x2 * scale2 / 2**(shift2-shift_reduce) + 0.5).floor()
    x = x1 + x2 + bias
    x = x / 2**shift_reduce
    x = (x + 0.5).floor()
    x = torch.clamp(x, quant_min, quant_max)
    return x

# def addquantize_function_per_tensor_int(x1, x2, scale1, shift1, scale2, shift2, bias, quant_min, quant_max, shift_reduce):
#     x1 = (x1.long() * scale1.long() + 2 ** (shift1-shift_reduce-1)).long() >> (shift1-shift_reduce).long()
#     x2 = (x2.long() * scale2.long() + 2 ** (shift2-shift_reduce-1)).long() >> (shift2-shift_reduce).long()
#     # x1 = (x1.long() * scale1.long()).long() >> (shift1-shift_reduce).long()
#     # x2 = (x2.long() * scale2.long()).long() >> (shift2-shift_reduce).long()
#     x = x1.long() + x2.long() + bias.long()
#     x = (x + 2**(shift_reduce - 1)).long().long() >> shift_reduce.long()
#     # x = (x).long().long() >> shift_reduce.long()
#     x = x.double()
#     x = torch.clamp(x, quant_min, quant_max)
#     return x.double()

# def addquantize_function_per_channel_int(x1, x2, scale1, shift1, scale2, shift2, bias, quant_min, quant_max, ch_axis, shift_reduce):
#     new_shape = [1] * len(x1.shape)
#     new_shape[ch_axis] = -1
#     scale1 = scale1.reshape(new_shape)
#     shift1 = shift1.reshape(new_shape)
#     scale2 = scale2.reshape(new_shape)
#     shift2 = shift2.reshape(new_shape)
#     bias = bias.reshape(new_shape)

#     x1 = (x1.long() * scale1.long() + 2 ** (shift1-shift_reduce-1)).long() >> (shift1-shift_reduce).long()
#     x2 = (x2.long() * scale2.long() + 2 ** (shift2-shift_reduce-1)).long() >> (shift2-shift_reduce).long()
#     # x1 = (x1.long() * scale1.long()).long() >> (shift1-shift_reduce).long()
#     # x2 = (x2.long() * scale2.long()).long() >> (shift2-shift_reduce).long()
#     x = x1.long() + x2.long() + bias.long()
#     x = (x + 2**(shift_reduce - 1)).long().long() >> shift_reduce.long()
#     # x = (x).long().long() >> shift_reduce.long()
#     x = x.double()
#     x = torch.clamp(x, quant_min, quant_max)
#     return x.double()

class AddQuantizePertensorInt(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x1, x2, scale1, shift1, scale2, shift2, bias, quant_min, quant_max, shift_reduce):
        return addquantize_function_per_tensor_int(x1, x2, scale1, shift1, scale2, shift2, bias, quant_min, quant_max, shift_reduce)

    @staticmethod
    def symbolic(g, x1, x2, scale1, shift1, scale2, shift2, bias, quant_min, quant_max, shift_reduce):
        return g.op("::AddQuantizePertensorInt", x1, x2, scale1, shift1, scale2, shift2, bias, shift_reduce, quant_min_i=quant_min, quant_max_i=quant_max)


class AddQuantizePerchannelInt(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x1, x2, scale1, shift1, scale2, shift2, bias, quant_min, quant_max, ch_axis, shift_reduce):
        return addquantize_function_per_channel_int(x1, x2, scale1, shift1, scale2, shift2, bias, quant_min, quant_max, ch_axis, shift_reduce)

    @staticmethod
    def symbolic(g, x1, x2, scale1, shift1, scale2, shift2, bias, quant_min, quant_max, ch_axis, shift_reduce):
        return g.op("::AddQuantizePerchannelInt", x1, x2, scale1, shift1, scale2, shift2, bias, shift_reduce, quant_min_i=quant_min, quant_max_i=quant_max, ch_axis_i=ch_axis)
