# modify from mqbench/observer.py
import warnings
import math
from functools import partial
from typing import Tuple
from copy import deepcopy
import torch
from torch.quantization.observer import _ObserverBase
from torch.quantization.fake_quantize import _is_per_channel, _is_per_tensor

from mqbench.utils import sync_tensor, pot_quantization, is_symmetric_quant
from mqbench.utils.logger import logger


class ObserverBaseCustom(_ObserverBase):
    '''
        Support per-tensor / per-channel.
        dtype: quant min/max can be infered using dtype, we actually do not need this.
        qscheme: quantization scheme
        reduce_range: special for fbgemm to avoid overflow
        quant_min: fix point value min
        quant_max: fix point value max
        ch_axis: per-channel axis or per-tensor(-1)
        above is similiar to torch observer.
        pot_scale: indecate wheather scale is power of two.
    '''

    min_val: torch.Tensor
    max_val: torch.Tensor

    def __init__(self, dtype=torch.quint8, qscheme=torch.per_tensor_affine, half_range=False,
                 reduce_range=False, bit=8, ch_axis=1, pot_scale=False, limit_min=None, limit_max=None, **kwargs):
        super(ObserverBaseCustom, self).__init__(dtype, qscheme, reduce_range, **kwargs)
        self.ch_axis = ch_axis
        self.pot_scale = pot_scale
        self.half_range = half_range
        self.bit = bit
        self.register_buffer("min_val", torch.tensor(float("inf")))
        self.register_buffer("max_val", torch.tensor(float("-inf")))
        self.quant_min, self.quant_max = self._calculate_qmin_qmax()
        self.limit_min = limit_min
        self.limit_max = limit_max

    @torch.jit.export
    def calculate_qparams(self):
        return self._calculate_qparams(self.min_val, self.max_val)

    def _calculate_qparams(self, min_val, max_val) -> Tuple[torch.Tensor, torch.Tensor]:
        r"""Calculates the quantization parameters, given min and max
        value tensors. Works for both per tensor and per channel cases

        Args:
            min_val: Minimum values per channel
            max_val: Maximum values per channel

        Returns:
            scales: Scales tensor of shape (#channels,)
            zero_points: Zero points tensor of shape (#channels,)
        """
        if min_val.numel() == 0 or max_val.numel() == 0:
            warnings.warn(
                "must run observer before calling calculate_qparams.\
                                    Returning default scale and zero point "
            )
            return torch.tensor([1.0]), torch.tensor([0])

        if min_val.dim() == 0 or max_val.dim() == 0:
            if min_val == float('inf') and self.max_val == float('-inf'):
                warnings.warn(
                    "must run observer before calling calculate_qparams.\
                                        Returning default scale and zero point "
                )
                return torch.tensor([1.0]), torch.tensor([0])

            assert min_val <= max_val, "min {} should be less than max {}".format(
                min_val, max_val
            )
        else:
            assert torch.all(min_val <= max_val), "min {} should be less than max {}".format(
                min_val, max_val
            )

        min_val, max_val = self.limit_min_max(min_val, max_val)

        mask = (min_val == max_val) * (max_val == 0)
        masked_min = torch.zeros_like(min_val)
        masked_max = torch.ones_like(max_val)
        if not self.half_range:
            masked_min -= 0.5
            masked_max -= 0.5

        min_val = min_val + masked_min * mask
        max_val = max_val + masked_max * mask

        quant_min, quant_max = self.quant_min, self.quant_max
        min_val_neg = torch.min(min_val, torch.zeros_like(min_val))
        max_val_pos = torch.max(max_val, torch.zeros_like(max_val))

        device = min_val_neg.device
        scale = torch.ones(min_val_neg.size(), dtype=torch.float32, device=device)
        zero_point = torch.zeros(min_val_neg.size(), dtype=torch.float32, device=device)

        if self.qscheme == torch.per_tensor_symmetric or self.qscheme == torch.per_channel_symmetric:
            max_val_pos = torch.max(-min_val_neg, max_val_pos)
            scale = max_val_pos / float(quant_max)
            scale = torch.max(scale, self.eps)
            
        elif self.qscheme == torch.per_channel_affine_float_qparams:
            scale = (max_val - min_val) / float(quant_max - quant_min)
            scale = torch.where(scale > self.eps, scale, torch.ones_like(scale))
            # We use the quantize function
            # xq = Round(Xf * inv_scale + zero_point),
            # setting zero_point to (-1 * min *inv_scale) we get
            # Xq = Round((Xf - min) * inv_scale)
            zero_point = -1 * min_val / scale
        else:
            scale = (max_val_pos - min_val_neg) / float(quant_max - quant_min)
            scale = torch.max(scale, self.eps)
            zero_point = quant_min - torch.round(min_val_neg / scale)
            zero_point = torch.clamp(zero_point, quant_min, quant_max)

        # For scalar values, cast them to Tensors of size 1 to keep the shape
        # consistent with default values in FakeQuantize.
        if len(scale.shape) == 0:
            # TODO: switch to scale.item() after adding JIT support
            scale = torch.tensor([float(scale)], dtype=scale.dtype, device=device)
        if len(zero_point.shape) == 0:
            # TODO: switch to zero_point.item() after adding JIT support
            zero_point = torch.tensor([int(zero_point)], dtype=zero_point.dtype, device=device)
            if self.qscheme == torch.per_channel_affine_float_qparams:
                zero_point = torch.tensor([float(zero_point)], dtype=zero_point.dtype, device=device)
        if self.pot_scale:
            scale = pot_quantization(scale)

        scale = sync_tensor(scale)
        zero_point = sync_tensor(zero_point)
      
        return scale, zero_point

    @torch.jit.export
    def _calculate_qmin_qmax(self) -> Tuple[int, int]:
        r"""Calculates actual qmin and qmax based on the quantization range,
        observer datatype and if range is reduced.
        """
        quant_min = -2 ** (self.bit - 1) if not self.half_range else 0
        quant_max = 2 ** (self.bit - 1) - 1 if not self.half_range else 2 ** self.bit - 1

        if self.reduce_range:
            quant_min, quant_max = quant_min // 2, quant_max // 2

        return quant_min, quant_max

    @torch.jit.export
    def extra_repr(self):
        return "min_val={}, max_val={}, pot={}".format(self.min_val if not _is_per_channel(self.qscheme) else 'List',
                                                                 self.max_val if not _is_per_channel(self.qscheme) else 'List',
                                                                 self.pot_scale)

    @torch.jit.export
    def limit_min_max(self, min_val:torch.Tensor, max_val:torch.Tensor):
        if self.limit_min is not None:
            min_val.clamp_(min=self.limit_min)

        if self.limit_max is not None:
            max_val.clamp_(max=self.limit_max)

        self.min_val = min_val
        self.max_val = max_val

        return min_val, max_val


class PseudoObserver(ObserverBaseCustom):
    '''
    Pseudo observer with fixed min_val and max_val.
    '''

    def __init__(self, **kwargs):

        super(PseudoObserver, self).__init__(**kwargs)

        if self.limit_min is not None:
            min_val = self.limit_min
        else:
            if self.half_range:
                min_val = 0
            else:
                min_val = -0.5

        if self.limit_max is not None:
            max_val = self.limit_max
        else:
            if self.half_range:
                max_val = 1
            else:
                max_val = 0.5
            
        assert min_val < max_val, 'max_val must bigger than min_val'
        self.min_val.fill_(min_val)
        self.max_val.fill_(max_val)


    def forward(self, x):
        return x
    

class MinMaxObserver(ObserverBaseCustom):
    '''
    Calculate minmax of whole calibration dataset.
    '''

    def __init__(self, **kwargs):
        super(MinMaxObserver, self).__init__(**kwargs)

    def forward(self, x_orig):
        r"""Records the running minimum and maximum of ``x``."""
        if x_orig.numel() == 0:
            return x_orig
        x = x_orig.to(self.min_val.dtype)
        if _is_per_tensor(self.qscheme):
            min_val_cur, max_val_cur = torch._aminmax(x)
        else:
            x_dim = x.size()
            new_axis_list = [i for i in range(len(x_dim))]
            new_axis_list[self.ch_axis] = 0
            new_axis_list[0] = self.ch_axis
            y = x.permute(new_axis_list)
            y = torch.flatten(y, start_dim=1)
            min_val_cur, max_val_cur = torch._aminmax(y, 1)
        self.min_val = torch.min(self.min_val, min_val_cur)
        self.max_val = torch.max(self.max_val, max_val_cur)

        return x
    

class EMAMinMaxObserver(ObserverBaseCustom):
    """Moving average min/max among batches.
    """

    def __init__(self, ema_ratio=0.9, **kwargs):
        super(EMAMinMaxObserver, self).__init__(**kwargs)
        self.ema_ratio = ema_ratio

    def forward(self, x_orig):
        r"""Records the running minimum and maximum of ``x``."""
        if x_orig.numel() == 0:
            return x_orig
        x = x_orig.to(self.min_val.dtype)
        if _is_per_tensor(self.qscheme):
            min_val_cur, max_val_cur = torch._aminmax(x)
        else:
            x_dim = x.size()
            new_axis_list = [i for i in range(len(x_dim))]  # noqa: C416
            new_axis_list[self.ch_axis] = 0
            new_axis_list[0] = self.ch_axis
            y = x.permute(new_axis_list)
            y = torch.flatten(y, start_dim=1)
            min_val_cur, max_val_cur = torch._aminmax(y, 1)

        if self.max_val.numel() <= 1 and self.max_val.isinf():
            self.min_val = min_val_cur
            self.max_val = max_val_cur
        else:
            self.min_val = self.min_val * self.ema_ratio + min_val_cur * (1.0 - self.ema_ratio)
            self.max_val = self.max_val * self.ema_ratio + max_val_cur * (1.0 - self.ema_ratio)
        return x