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


class MixedQuantizerInt(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,
                 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=0,
                 half_range=True, # if scale use int or uint
                 shift_bit=6,
                 **kwargs):
        super(MixedQuantizerInt, 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.shift_reduce = torch.nn.Parameter(torch.tensor([shift_reduce], device=scale.device, dtype=torch.float32), requires_grad=False)
        self.half_range = half_range
        self.per_channel_shift = kwargs.get('per_channel_shift', False)
        int_scale, shift, bias = self._get_fixpoint_scale_shift(scale, bias)
        self.scale = torch.nn.Parameter(int_scale.to(torch.float32), requires_grad=False)
        self.shift = torch.nn.Parameter(shift.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)

        self.clip_value_up = torch.nn.Parameter(torch.tensor([quant_max], dtype=torch.float32), requires_grad=False)
        self.clip_value_low = torch.nn.Parameter(torch.tensor([quant_min], dtype=torch.float32), requires_grad=False)
        
    def _get_fixpoint_scale_shift(self, scale, bias):
        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, bias.to(scale.device)

    @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={}, bias={}, shift={}, clip_value_low={}, clip_value_up={}, ' \
               '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.scale, self.bias,
                   self.shift, self.clip_value_low, self.clip_value_up, self.shift_reduce)

    def forward(self, x):
        if self.is_per_channel:
            if is_tracing_state():
                x = MixedQuantizePerchannelInt.apply(
                    x, self.scale, self.shift, self.bias, self.quant_min, self.quant_max, self.ch_axis, self.clip_value_low, self.clip_value_up, self.shift_reduce)
            else:
                x = mixedquantize_function_per_channel_int(
                    x, self.scale, self.shift, self.bias, self.quant_min, self.quant_max, self.ch_axis, self.clip_value_low, self.clip_value_up, self.shift_reduce)
        else:
            if is_tracing_state():
                x = MixedQuantizePertensorInt.apply(
                    x, self.scale, self.shift, self.bias, self.quant_min, self.quant_max, self.clip_value_low, self.clip_value_up, self.shift_reduce)
            else:
                x = mixedquantize_function_per_tensor_int(
                    x, self.scale, self.shift, self.bias, self.quant_min, self.quant_max, self.clip_value_low, self.clip_value_up, self.shift_reduce)
        
        return x

def mixedquantize_function_per_tensor_int(x, scale, shift, bias, quant_min, quant_max, clip_value_low, clip_value_up, shift_reduce):
    x = (x * scale / 2**(shift-shift_reduce) + 0.5).floor()
    x = x + bias
    x = (x / 2**shift_reduce + 0.5).floor()
    if clip_value_up is not None:
        x = torch.clamp(x, max=clip_value_up)
    if clip_value_low is not None:
        x = torch.clamp(x, min=clip_value_low)
    x = torch.clamp(x, quant_min, quant_max)
    return x


def mixedquantize_function_per_channel_int(x, scale, shift, bias, quant_min, quant_max, ch_axis, clip_value_low, clip_value_up, shift_reduce):
    new_shape = [1] * len(x.shape)
    new_shape[ch_axis] = -1
    scale = scale.reshape(new_shape)
    shift = shift.reshape(new_shape)
    bias = bias.reshape(new_shape)

    # x = (x * scale / 2**shift + 0.5).floor()
    x = (x * scale / 2**(shift-shift_reduce) + 0.5).floor()
    x = x + bias
    x = (x / 2**shift_reduce + 0.5).floor()
    if clip_value_up is not None:
        x = torch.clamp(x, max=clip_value_up)
    if clip_value_low is not None:
        x = torch.clamp(x, min=clip_value_low)
    x = torch.clamp(x, quant_min, quant_max)
    return x

# def mixedquantize_function_per_tensor_int(x, scale, shift, bias, quant_min, quant_max, clip_value_low, clip_value_up, shift_reduce):
#     x = (x.long() * scale.long() + 2 ** (shift-1)).long() >> shift.long()
#     # x = (x.long() * scale.long()).long() >> shift.long()
#     x = x.double()
#     if clip_value_up is not None:
#         x = torch.clamp(x, max=clip_value_up)
#     if clip_value_low is not None:
#         x = torch.clamp(x, min=clip_value_low)
#     x = torch.clamp(x, quant_min, quant_max)
#     return x.double()

# def mixedquantize_function_per_channel_int(x, scale, shift, bias, quant_min, quant_max, ch_axis, clip_value_low, clip_value_up, shift_reduce):
#     new_shape = [1] * len(x.shape)
#     new_shape[ch_axis] = -1
#     scale = scale.reshape(new_shape)
#     shift = shift.reshape(new_shape)


#     x = (x.long() * scale.long() + 2 ** (shift-1)).long() >> shift.long()
#     # x = (x.long() * scale.long()).long() >> shift.long()
#     x = x.double()
#     if clip_value_up is not None:
#         x = torch.clamp(x, max=clip_value_up)
#     if clip_value_low is not None:
#         x = torch.clamp(x, min=clip_value_low)
#     x = torch.clamp(x, quant_min, quant_max)
#     return x


class MixedQuantizePertensorInt(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale, shift, bias, quant_min, quant_max, clip_value_low, clip_value_up, shift_reduce):
        return mixedquantize_function_per_tensor_int(x, scale, shift, bias, quant_min, quant_max, clip_value_low, clip_value_up, shift_reduce)

    @staticmethod
    def symbolic(g, x, scale, shift, bias, quant_min, quant_max, clip_value_low, clip_value_up, shift_reduce):
        return g.op("::MixedQuantizePertensorInt", x, scale, shift, bias, clip_value_low, clip_value_up, shift_reduce, quant_min_i=quant_min, quant_max_i=quant_max)


class MixedQuantizePerchannelInt(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale, shift, bias, quant_min, quant_max, ch_axis, clip_value_low, clip_value_up, shift_reduce):
        return mixedquantize_function_per_channel_int(x, scale, shift, bias, quant_min, quant_max, ch_axis, clip_value_low, clip_value_up, shift_reduce)

    @staticmethod
    def symbolic(g, x, scale, shift, bias, quant_min, quant_max, ch_axis, clip_value_low, clip_value_up, shift_reduce):
        return g.op("::MixedQuantizePerchannelInt", x, scale, shift, bias, clip_value_low, clip_value_up, shift_reduce, quant_min_i=quant_min, quant_max_i=quant_max, ch_axis_i=ch_axis)