# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""
Classic uniform quantization over n bits.
"""
from typing import Tuple
import torch

from .base import BaseQuantizer
from .utils import simple_repr


def uniform_quantize(p: torch.Tensor, bits: torch.Tensor = torch.tensor(8.)):
    """
    Quantize the given weights over `bits` bits.

    Returns:
        - quantized levels
        - (min, max) range.

    """
    assert (bits >= 1).all() and (bits <= 15).all()
    num_levels = (2 ** bits.float()).long()
    mn = p.min().item()
    mx = p.max().item()
    p = (p - mn) / (mx - mn)  # put p in [0, 1]
    unit = 1 / (num_levels - 1)  # quantization unit
    levels = (p / unit).round()
    if (bits <= 8).all():
        levels = levels.byte()
    else:
        levels = levels.short()
    return levels, (mn, mx)


def uniform_unquantize(levels: torch.Tensor, scales: Tuple[float, float],
                       bits: torch.Tensor = torch.tensor(8.)):
    """
    Unquantize the weights from the levels and scale. Return a float32 tensor.
    """
    mn, mx = scales
    num_levels = 2 ** bits.float()
    unit = 1 / (num_levels - 1)
    levels = levels.float()
    p = levels * unit  # in [0, 1]
    return p * (mx - mn) + mn


class UniformQuantizer(BaseQuantizer):
    def __init__(self, model: torch.nn.Module, bits: float = 8., min_size: float = 0.01,
                 float16: bool = False, qat: bool = False, exclude=[], detect_bound=True):
        """
        Args:
            model (torch.nn.Module): model to quantize
            bits (float): number of bits to quantize over.
            min_size (float): minimum size in MB of a parameter to be quantized.
            float16 (bool): if a layer is smaller than min_size, should we still do float16?
            qat (bool): perform quantized aware training.
            exclude (list[str]): list of patterns used to match parameters to exclude.
                For instance `['bias']` to exclude all bias terms.
            detect_bound (bool): if True, will detect bound parameters and reuse
                the same quantized tensor for both.
        """
        self.bits = float(bits)
        self.qat = qat

        super().__init__(model, min_size, float16, exclude, detect_bound)

    def __repr__(self):
        return simple_repr(self, )

    def _pre_forward_train(self):
        if self.qat:
            for qparam in self._qparams:
                if qparam.other is not None:
                    new_param = qparam.other.module._parameters[qparam.other.name]
                else:
                    quantized = self._quantize_param(qparam)
                    qvalue = self._unquantize_param(qparam, quantized)
                    new_param = qparam.param + (qvalue - qparam.param).detach()
                qparam.module._parameters[qparam.name] = new_param
            return True
        return False

    def _post_forward_train(self):
        if self.qat:
            for qparam in self._qparams:
                qparam.module._parameters[qparam.name] = qparam.param
            return True
        return False

    def _quantize_param(self, qparam):
        levels, scales = uniform_quantize(qparam.param.data, torch.tensor(self.bits))
        return (levels, scales)

    def _unquantize_param(self, qparam, quantized):
        levels, scales = quantized
        return uniform_unquantize(levels, scales, torch.tensor(self.bits))

    def model_size(self):
        """
        Non differentiable model size in MB.
        """
        total = super().model_size()
        subtotal = 0
        for qparam in self._qparams:
            if qparam.other is None:  # if parameter is bound, count only one copy.
                subtotal += self.bits * qparam.param.numel() + 64  # 2 float for the overall scales
        subtotal /= 2**20 * 8  # bits to MegaBytes
        return total + subtotal

    def true_model_size(self):
        """
        Return the true quantized model size, in MB, without extra
        compression.
        """
        return self.model_size().item()