# modify from 3rdparty/mqbench/fake_quantize/quantize_base.py
import torch
from torch.quantization import FakeQuantizeBase
from torch.quantization.observer import MovingAverageMinMaxObserver
from torch.quantization.fake_quantize import _is_per_channel, _is_per_tensor

from mqbench.utils import is_symmetric_quant

class QuantizeBase(FakeQuantizeBase):
    r""" This is an extension of the FakeQuantize module in fake_quantize.py, which
    supports more generalized lower-bit quantization and support learning of the scale
    and zero point parameters through backpropagation. For literature references,
    please see the class _LearnableFakeQuantizePerTensorOp.
    In addition to the attributes in the original FakeQuantize module, the _LearnableFakeQuantize
    module also includes the following attributes to support quantization parameter learning.
    """
    _per_channel_state_dict = ['scale', 'zero_point']

    def __init__(self, observer=MovingAverageMinMaxObserver, **observer_kwargs):
        super().__init__()
        self.observer_kwargs = observer_kwargs
        self.activation_post_process = observer(**observer_kwargs)
        self.dtype = self.activation_post_process.dtype
        self.qscheme = self.activation_post_process.qscheme
        self.quant_min = self.activation_post_process.quant_min
        self.quant_max = self.activation_post_process.quant_max
        self.half_range = self.activation_post_process.half_range
        bitwidth = self.activation_post_process.bit
        assert self.quant_min <= self.quant_max, \
            'quant_min must be less than or equal to quant_max'
        self.pot_scale = self.activation_post_process.pot_scale
        self.ch_axis = self.activation_post_process.ch_axis \
            if hasattr(self.activation_post_process, 'ch_axis') else 1
        assert _is_per_channel(self.qscheme) or \
            _is_per_tensor(self.qscheme), \
            'Only per channel and per tensor quantization are supported in fake quantize' + \
            ' got qscheme: ' + str(self.qscheme)
        self.is_per_channel = _is_per_channel(self.qscheme)
        self.is_symmetric_quant = is_symmetric_quant(self.qscheme)
        
        self.ckpt_bitwidth = 32
        self.register_buffer('is_initialized_params', torch.tensor([0], dtype=torch.uint8))
        self.register_buffer('bitwidth', torch.tensor(bitwidth, dtype=torch.uint8, requires_grad=False))


    def forward(self, x):
        x = self._forward_observe(x)
        x = self._forward_quant(x)
        return x

    def _forward_quant(self, X):
        raise NotImplementedError('Fake Quantize must implement _forward_quant method')

    def _forward_observe(self, X):
        if self.observer_enabled[0] == 1:
            self.activation_post_process(X.detach())
        return X

    def init_qparams(self):
        if self.is_initialized_params[0] == 0:
            _scale, _zero_point = self.calculate_qparams()
            _scale, _zero_point = _scale.to(self.scale.device), _zero_point.to(self.zero_point.device)
            if self.is_per_channel:
                self.scale.data = torch.ones_like(_scale)
                self.zero_point.data = torch.zeros_like(_zero_point.float())
            self.scale.data.copy_(_scale)
            self.zero_point.data.copy_(_zero_point.float())
            self.set_initialized()

    @torch.jit.export
    def calculate_qparams(self):
        return self.activation_post_process.calculate_qparams()

    @torch.jit.export
    def extra_repr(self):
        return 'fake_quant_enabled={}, observer_enabled={}, is_initialized_params={}, ' \
               'bitwidth={}, quant_min={}, quant_max={}, dtype={}, qscheme={}, ch_axis={}, is_per_channel={}, is_symmetric_quant={}'.format(
                   self.fake_quant_enabled, self.observer_enabled, self.is_initialized_params,
                   self.bitwidth, self.quant_min, self.quant_max,
                   self.dtype, self.qscheme, self.ch_axis, self.is_per_channel, self.is_symmetric_quant)

    @torch.jit.export
    def to_quantizer_params(self):
        return {
            'quant_min': self.quant_min,
            'quant_max': self.quant_max,
            'scale': self.scale if hasattr(self, 'scale') else None,
            'zero_point': self.zero_point.data.float() if hasattr(self, 'zero_point') else None,
            'ch_axis': self.ch_axis,
            'is_per_channel': self.is_per_channel,
            'is_symmetric_quant': self.is_symmetric_quant,
            'bitwidth': self.bitwidth
        }

    @torch.jit.export
    def set_initialized(self, initialized: bool = True) -> None:
        self.is_initialized_params[0] = 1 if initialized else 0

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
                              missing_keys, unexpected_keys, error_msgs):
        local_state = self._per_channel_state_dict
        for name in local_state:
            key = prefix + name
            if key in state_dict:
                val = state_dict[key]
                attr = getattr(self, name, None)
                if attr is not None and attr.shape != val.shape:
                    attr.data = torch.ones_like(val, device=attr.data.device)
        
        super(QuantizeBase, self)._load_from_state_dict(state_dict, prefix, local_metadata, strict,
                                                        missing_keys, unexpected_keys, error_msgs)
