#
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#


"""TensorQuantizer Module"""
import math
from absl import logging

import torch
from torch import nn

from pytorch_quantization.tensor_quant import QuantDescriptor, tensor_quant, fake_tensor_quant
from pytorch_quantization.nn.modules.clip import Clip

from pytorch_quantization import calib

import pytorch_quantization.utils as quant_utils

__all__ = ['TensorQuantizer']

class TensorQuantizer(nn.Module):
    """Tensor quantizer module

    This module uses tensor_quant or fake_tensor_quant function to quantize a tensor. And wrappers variable, moving
    statistics we'd want when training a quantized network.

    Experimental features:
        ``clip`` stage learns range before enabling quantization.
        ``calib`` stage runs calibration

    Args:
        quant_desc: An instance of :func:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
        disabled: A boolean. If True, by pass the whole module returns input. Default False.
        if_quant: A boolean. If True, run main quantization body. Default True.
        if_clip: A boolean. If True, clip before quantization and learn amax. Default False.
        if_calib: A boolean. If True, run calibration. Not implemented yet. Settings of calibration will probably
            go to :func:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.

    Raises:

    Readonly Properties:
        - axis:
        - fake_quant:
        - scale:
        - step_size:

    Mutable Properties:
        - num_bits:
        - unsigned:
        - amax:
    """

    # An experimental static switch for using pytorch's native fake quantization
    # Primary usage is to export to ONNX
    use_fb_fake_quant = False

    def __init__(self, quant_desc=QuantDescriptor(), disabled=False, if_quant=True, if_clip=False, if_calib=False):
        """Initialize quantizer and set up required variables"""
        super(TensorQuantizer, self).__init__()
        # Expand quant_desc. Use quant_desc.dict would be eaiser, but adding one-by-one explicitly gives more control
        self._num_bits = quant_desc.num_bits
        self._fake_quant = quant_desc.fake_quant
        self._axis = quant_desc.axis
        self._scale_amax = quant_desc.scale_amax
        self._learn_amax = quant_desc.learn_amax
        self._unsigned = quant_desc.unsigned
        self._narrow_range = quant_desc.narrow_range

        self._scale = None if not quant_desc.fake_quant else 1.
        self._disabled = disabled
        self._if_quant = if_quant
        self._if_clip = False
        self._if_calib = if_calib

        if quant_desc.amax is not None:
            self.register_buffer('_amax', torch.tensor(quant_desc.amax))

        # Clip module consumes a lot of memory, so only create it if learn_amax is True
        if self._learn_amax:
            init_amax = quant_desc.amax if quant_desc.amax is not None else 1.
            self.clip = Clip(-init_amax, init_amax, learn_min=True, learn_max=True)
            # It makes more sense to enable clip stage (which learns amax) if learn_amax is true
            self.enable_clip()
        if if_clip:
            self.enable_clip()

        if quant_desc.calib_method == "histogram":
            logging.info("Creating histogram calibrator")
            self._calibrator = calib.HistogramCalibrator(
                num_bits=self._num_bits, axis=self._axis, unsigned=self._unsigned)
        elif quant_desc.calib_method == "max":
            logging.info("Creating Max calibrator")
            self._calibrator = calib.MaxCalibrator(num_bits=self._num_bits, axis=self._axis, unsigned=self._unsigned)

    # pylint:disable=missing-docstring
    @property
    def num_bits(self):
        return self._num_bits

    @property
    def unsigned(self):
        return self._unsigned

    @property
    def scale(self):
        if self._fake_quant:
            logging.error("Fake quantize mode doesn't use scale explicitly!")
        if self._scale is None:
            logging.critical("Accessing scale before quantizing any tensor!")
        return self._scale

    @property
    def amax(self):
        if not hasattr(self, "_amax"):
            return None
        return self._amax

    @property
    def step_size(self):
        if not hasattr(self, "_amax"):
            logging.error("step_size is undefined under dynamic amax mode!")
            return None
        return self._amax / (2.0**(self._num_bits - 1 + int(self._unsigned)) - 1.0)

    @property
    def axis(self):
        return self._axis

    @property
    def fake_quant(self):
        return self._fake_quant

    @property
    def narrow_range(self):
        return self._narrow_range

    def disable(self):
        """Bypass the module"""
        self._disabled = True

    def enable(self):
        self._disabled = False

    def disable_clip(self):
        """Disable clip stage"""
        self._if_clip = False
        self.clip.clip_value_min.required_grad = False
        self.clip.clip_value_max.required_grad = False

    def enable_clip(self):
        """Enable clip stage"""
        logging.warning("Enable `clip` stage for amax learning.")
        if not self._learn_amax:
            raise ValueError("learn_amax is False. Cannot enable clip.")
        self.clip.clip_value_min.required_grad = True
        self.clip.clip_value_max.required_grad = True
        self._if_clip = True

    def disable_calib(self):
        logging.warning("Disable {}".format(self._calibrator.__class__.__name__))
        self._if_calib = False

    def enable_calib(self):
        if self._calibrator is None:
            raise ValueError("Calibrator was not created, cannot enable calibration.")
        logging.info("Enable {}".format(self._calibrator.__class__.__name__))
        self._if_calib = True

    def disable_quant(self):
        logging.info("Disable `quant` stage.")
        self._if_quant = False

    def enable_quant(self):
        logging.info("Enable `quant` stage.")
        self._if_quant = True

    @amax.setter
    def amax(self, value):
        if value is None:
            logging.error("Setting amax no None is meaningless.")
        else:
            if isinstance(value, torch.Tensor):
                logging.warning("amax setter is not designed to take tensor.")
            if not hasattr(self, "_amax"):
                self.register_buffer('_amax', torch.tensor(value))
            else:
                value = torch.tensor(value, device=self._amax.device)
                if self._amax.shape != value.shape:
                    raise TypeError("Changing shape when setting amax is not allowed.")
                self._amax.data.copy_(value.data)

    @num_bits.setter
    def num_bits(self, value):
        self._num_bits = value

    @unsigned.setter
    def unsigned(self, value):
        self._unsigned = value

    @narrow_range.setter
    def narrow_range(self, value):
        self._narrow_range = value

    # pylint:enable=missing-docstring
    def load_calib_amax(self, *args, **kwargs):
        """Load amax from calibrator.

        Updates the amax buffer with value computed by the calibrator, creating it if necessary.
        *args and **kwargs are directly passed to compute_amax, except "strict" in kwargs. Refer to
        compute_amax for more details.
        """
        strict = kwargs.pop("strict", True)
        if getattr(self, '_calibrator', None) is None:
            raise RuntimeError("Calibrator not created.")
        calib_amax = self._calibrator.compute_amax(*args, **kwargs)
        if calib_amax is None:
            err_msg = "Calibrator returned None. This usually happens when calibrator hasn't seen any tensor."
            if not strict:
                logging.warning(err_msg)
                logging.warning("Set amax to NaN!")
                calib_amax = torch.tensor(math.nan)
            else:
                raise RuntimeError(err_msg + " Passing 'strict=False' to `load_calib_amax()` will ignore the error.")
        logging.warning("Load calibrated amax, shape={}.".format(calib_amax.shape))
        logging.log_first_n(
            logging.WARNING, "Call .cuda() if running on GPU after loading calibrated amax.", 1)
        if not hasattr(self, '_amax'):
            self.register_buffer('_amax', calib_amax.data)
        else:
            self._amax.copy_(calib_amax)

    def init_learn_amax(self):
        """Initialize learned amax from fixed amax"""
        if self._learn_amax is False:
            raise RuntimeError("Called init_learn_amax with learn_amax=False.")
        logging.warning("Load amax as initial value for amax learning!")
        if self._amax.numel() != 1:
            logging.warning("Per channel learned amax not supported. Initializing with max(amax).")
            init_amax = torch.max(self._amax)
        else:
            init_amax = self._amax
        self.clip.clip_value_min.data.copy_(-init_amax.data)
        self.clip.clip_value_max.data.copy_(init_amax.data)

    def _get_amax(self, inputs):
        """get amax from buffer or compute it dynamically."""
        if hasattr(self, '_amax'):
            amax = self._amax
        else:
            if self._axis is None:
                reduce_axis = None
            else:
                reduce_axis = []
                # Swap axis to reduce
                axis = self._axis if isinstance(self._axis, (list, tuple)) else [self._axis]
                for i in range(inputs.dim()):
                    if not i in axis:
                        reduce_axis.append(i)
            amax = quant_utils.reduce_amax(inputs, axis=reduce_axis, keepdims=True).detach()
        if self._scale_amax is not None:
            amax = amax.detach() * self._scale_amax

        return amax

    def _fb_fake_quant(self, inputs, amax):
        """Native pytorch fake quantization."""
        logging.log_first_n(logging.WARNING, "Use Pytorch's native experimental fake quantization.", 1)
        bound = (1 << (self._num_bits - 1 + int(self._unsigned))) - 1
        # To be consistent with ONNX, full range is used. e.g. range is [-128, 127] in int8
        if amax.numel() == 1:
            outputs = torch.fake_quantize_per_tensor_affine(
                inputs, amax.item() / bound, 0,
                -bound - 1 if not self._unsigned else 0, bound)
        else:
            amax_sequeeze = amax.squeeze().detach()
            if len(amax_sequeeze.shape) != 1:
                raise TypeError("Pytorch's native quantization doesn't support multiple axes")
            quant_dim = list(amax.shape).index(list(amax_sequeeze.shape)[0])
            scale = amax_sequeeze / bound
            outputs = torch.fake_quantize_per_channel_affine(
                inputs, scale.data, torch.zeros_like(scale, dtype=torch.long).data, quant_dim,
                -bound - 1 if not self._unsigned else 0, bound)

        return outputs

    def _quant_forward(self, inputs):
        """Quantized forward pass."""
        if self._learn_amax:
            inputs = self.clip(inputs)
            amax = torch.max(-self.clip.clip_value_min, self.clip.clip_value_max).detach()
        else:
            amax = self._get_amax(inputs)

        if self._fake_quant:
            if not TensorQuantizer.use_fb_fake_quant:
                outputs = fake_tensor_quant(inputs, amax, self._num_bits, self._unsigned, self._narrow_range)
            else:
                if inputs.dtype == torch.half or amax.dtype == torch.half:
                    raise Exception("Exporting to ONNX in fp16 is not supported. Please export in fp32, i.e. disable AMP.")
                outputs = self._fb_fake_quant(inputs, amax)
        else:
            outputs, self._scale = tensor_quant(inputs, amax, self._num_bits, self._unsigned)

        return outputs

    def forward(self, inputs):
        """Apply tensor_quant function to inputs

        Args:
            inputs: A Tensor of type float32.

        Returns:
            outputs: A Tensor of type output_dtype
        """
        if self._disabled:
            return inputs

        outputs = inputs

        if self._if_calib:
            if self._calibrator is None:
                raise RuntimeError("Calibrator was not created.")
            # Shape is only know when it sees the first tensor
            self._calibrator.collect(inputs)

        if self._if_clip:
            if not self._learn_amax:
                raise RuntimeError("Clip without learning amax is not implemented.")
            outputs = self.clip(inputs)

        if self._if_quant:
            outputs = self._quant_forward(inputs)

        return outputs

    def _short_amax(self, fmt='.4f'):
        """Short description of amax

        Returns:
            'dynamic': if _amax is not registered
            'amax': if _amax is per-tensor
            '[min, max](size)': if _amax is per-channel
        """
        if not hasattr(self, '_amax'):
            return 'dynamic'
        if self._amax.numel() == 1:
            return '{:{fmt}}'.format(self._amax.item(), fmt=fmt)
        return '[{:{fmt}}, {:{fmt}}]({})'.format(self._amax.min().item(), self._amax.max().item(),
                                                 self._amax.numel(), fmt=fmt)

    def extra_repr(self):
        if self._disabled:
            return "disabled"
        s = "{}{}bit".format("unsigned " if self._unsigned else "", self._num_bits)
        s += " narrow" if (self._narrow_range) else ""
        s += " fake" if (self._fake_quant) else ""
        s += " axis={}".format(self._axis) if self._axis is not None else " per-tensor"
        s += " amax={}".format(self._short_amax())
        s += " *{}".format(self._scale_amax) if self._scale_amax else ""
        s += " learned" if (self._learn_amax) else ""
        s += " calibrator={}".format(self._calibrator.__class__.__name__) if (self._calibrator is not None) else ""
        s += " scale={}".format(self._scale) if self._scale is not None else ""
        s += " quant" if (self._if_quant) else ""
        s += " clip" if (self._if_clip) else ""
        s += " calib" if (self._if_calib) else ""
        return s

    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
        """Overloaded module function

        Adds warnings during state_dict loading.
        A workaround is implemented for loading amax from checkpoint and only supports CUDA.

        Args:
            state_dict: A dict containing the state of the top level module
            prefix: A string that prefixes all of this modules state in state_dict, e.g. 'model.conv1.'
        """
        dst_has_amax = '_amax' in self._buffers
        src_has_amax = prefix + '_amax' in state_dict

        if not src_has_amax and dst_has_amax:
            logging.error("{}: No amax in state_dict.".format(prefix[:-1]))
        elif src_has_amax and not dst_has_amax:
            logging.debug(("{}: No '_amax' buffer to load amax into."
                           " '_amax` will be created as WAR for now. "
                           "This behavior will change in future.").format(prefix[:-1]))
            self.register_buffer("_amax", state_dict[prefix + '_amax'].data.cuda())
        elif src_has_amax and dst_has_amax:
            logging.warning("{}: Overwriting amax.".format(prefix[:-1]))

        super(TensorQuantizer, self)._load_from_state_dict(state_dict, prefix, *args, **kwargs)
