#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
#


import tensorflow as tf
from tensorflow.python.util import tf_inspect
from tensorflow_quantization.quantize_wrapper_base import BaseQuantizeWrapper
import warnings

"""
Naming convention for keras `layer` quantize wrapper is
<layer.__class__.__name__>QuantizeWrapper
"""

DISABLED_LAYER_QUANTIZATION_DEFAULT = [
    "MaxPooling2D",
    "BatchNormalization",
    "Add",
    "Multiply",
    "Concatenate"
]


# ##############################################
# ############# Weighted Layers ################
# ##############################################


class WeightedBaseQuantizeWrapper(BaseQuantizeWrapper):
    """
    BaseQuantizeWrapper for weighted layers: Conv2D, DepthwiseConv2D, and Dense layer.
        These layers share a lot of the same code except for a few modifications. Conv2D and Dense share the same code.
        Layers that inherit this class support weight and input QDQ nodes.

    TRT Rule:
        One Q/DQ pair should be attached to the input activation, and another Q/DQ pair should be attached to weights.
        Weights tensor is per-channel quantized:
            For the Q/DQ attached to weight tensor, set axis=0 and axis=1 for Conv and ConvTransposed respectively.
        Input tensor is per-tensor quantized.
    """

    def __init__(
        self, layer: tf.keras.layers.Layer, kernel_type: str = "kernel", **kwargs
    ):
        """
        Creates a wrapper to emulate quantization for a keras layer.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          kernel_type (str): Options=['kernel' for Conv2D/Dense, 'depthwise_kernel' for DepthwiseConv2D]
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        self.kernel_type = kernel_type
        self.channel_axis = kwargs.get("axis", -1)
        super().__init__(layer, **kwargs)

    def build(self, input_shape):
        super().build(input_shape)

        self._weight_vars = []
        self.input_vars = {}
        self.output_vars = {}
        self.channel_axis = -1
        if self.kernel_type == "depthwise_kernel":
            self.channel_axis = 2
        # quantize weights only applicable for weighted ops.
        # By default weights is per channel quantization
        if self.quantize_weights:
            # get kernel weights dims.
            kernel_weights = getattr(self.layer, self.kernel_type)
            min_weight = self.layer.add_weight(
                kernel_weights.name.split(":")[0] + "_min",
                shape=(kernel_weights.shape[self.channel_axis]),
                initializer=tf.keras.initializers.Constant(-6.0),
                trainable=False,
            )
            max_weight = self.layer.add_weight(
                kernel_weights.name.split(":")[0] + "_max",
                shape=(kernel_weights.shape[self.channel_axis]),
                initializer=tf.keras.initializers.Constant(6.0),
                trainable=False,
            )
            quantizer_vars = {"min_var": min_weight, "max_var": max_weight}
            self._weight_vars.append((kernel_weights, quantizer_vars))
            # Needed to ensure unquantized weights get trained as part of the wrapper.
            self._trainable_weights.append(kernel_weights)

        # By default input is per tensor quantization
        if self.quantize_inputs:
            input_min_weight = self.layer.add_weight(
                self.layer.name + "_ip_min",
                shape=None,
                initializer=tf.keras.initializers.Constant(-6.0),
                trainable=False,
            )
            input_max_weight = self.layer.add_weight(
                self.layer.name + "_ip_max",
                shape=None,
                initializer=tf.keras.initializers.Constant(6.0),
                trainable=False,
            )
            self.input_vars["min_var"] = input_min_weight
            self.input_vars["max_var"] = input_max_weight

    def call(self, inputs, training=None):
        if training is None:
            training = tf.keras.backend.learning_phase()

        # Quantize all weights, and replace them in the underlying layer.
        if self.quantize_weights:
            quantized_weights = []
            quantized_weight = self._last_value_quantizer(
                self._weight_vars[0][0],
                training,
                self._weight_vars[0][1],
                per_channel=True,
                channel_axis=self.channel_axis,
            )
            quantized_weights.append(quantized_weight)
            # Replace the original weights with QDQ weights
            setattr(self.layer, self.kernel_type, quantized_weights[0])

        # Quantize inputs to the conv layer
        if self.quantize_inputs:
            quantized_inputs = self._last_value_quantizer(
                inputs, training, self.input_vars, per_channel=False
            )
        else:
            quantized_inputs = inputs

        args = tf_inspect.getfullargspec(self.layer.call).args
        if "training" in args:
            outputs = self.layer.call(quantized_inputs, training=training)
        else:
            outputs = self.layer.call(quantized_inputs)

        return outputs


class Conv2DQuantizeWrapper(WeightedBaseQuantizeWrapper):
    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
        """
        Creates a wrapper to emulate quantization for the Conv2D keras layer.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        self.kernel_type = "kernel"
        super().__init__(layer, kernel_type=self.kernel_type, **kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, training=None):
        return super().call(inputs, training=training)


class DenseQuantizeWrapper(WeightedBaseQuantizeWrapper):
    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
        """
        Creates a wrapper to emulate quantization for the Dense keras layer.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        self.kernel_type = "kernel"
        super().__init__(layer, kernel_type=self.kernel_type, **kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, training=None):
        return super().call(inputs, training=training)


class DepthwiseConv2DQuantizeWrapper(WeightedBaseQuantizeWrapper):
    """Requires TF >= 2.8.0"""

    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
        """
        Creates a wrapper to emulate quantization for the DepthwiseConv2D keras layer.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        self.kernel_type = "depthwise_kernel"
        super().__init__(layer, kernel_type=self.kernel_type, **kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, training=None):
        return super().call(inputs, training=training)


# ##############################################
# ########### Non-Weighted Layers ##############
# ######### with Single Input/Output ###########
# ##############################################


class NonWeightedBaseQuantizeWrapper(BaseQuantizeWrapper):
    """
    BaseQuantizeWrapper for non-weighted layers with Single Input/Output: AveragePooling2D, GlobalAveragePooling,
        MaxPooling2D and BatchNormalization.

    Supports 1 input and 1 output QDQ. Similar to Concat, except that Concat supports multiple inputs.
    NonWeightedBaseQuantizeWrapper can use WeightedBaseQuantizeWrapper by giving quantize_weigths=False.
    """

    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
        """
        Creates a wrapper to emulate quantization for non-weighted keras layers.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        super().__init__(layer, **kwargs)

    def build(self, input_shape):
        super().build(input_shape)

        self.input_vars = {}

        # By default input is per tensor quantization
        if self.quantize_inputs:
            input_min_weight = self.layer.add_weight(
                self.layer.name + "_ip_min",
                shape=None,
                initializer=tf.keras.initializers.Constant(-6.0),
                trainable=False,
            )
            input_max_weight = self.layer.add_weight(
                self.layer.name + "_ip_max",
                shape=None,
                initializer=tf.keras.initializers.Constant(6.0),
                trainable=False,
            )
            self.input_vars["min_var"] = input_min_weight
            self.input_vars["max_var"] = input_max_weight

    def call(self, inputs, training=None):
        if training is None:
            training = tf.keras.backend.learning_phase()

        # Quantize inputs to the conv layer
        if self.quantize_inputs:
            quantized_inputs = self._last_value_quantizer(
                inputs, training, self.input_vars, per_channel=False
            )
        else:
            quantized_inputs = inputs

        args = tf_inspect.getfullargspec(self.layer.call).args
        if "training" in args:
            outputs = self.layer.call(quantized_inputs, training=training)
        else:
            outputs = self.layer.call(quantized_inputs)

        return outputs


class AveragePooling2DQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
    """
    TRT Rule:
        Add Q/DQ to its input if the ops follows is quantized.
        Quantize average pooling will introduce small variance compared to float because of the rounding change.
        TensorRT doesn’t have Int8 in and fp32 out average pool support.
        If the op follows average pooling is not quantized, it is users choice between running average pooling
        in int8 then convert to fp32 for the following op and run average pooling in fp32.

    Currently, we're adding QDQ to all AveragePooling2D layers.
    """

    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
        """
        Creates a wrapper to emulate quantization for the AveragePooling2D keras layer.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        super().__init__(layer, **kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, training=None):
        return super().call(inputs, training=training)


class GlobalAveragePooling2DQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
    """
    TRT Rule:
        No explicit rule from the TRT team. Following the same as AveragePooling2D.

    Residual block v2: Add to MaxPool (branch1) and BN (branch2).

    Supports 1 input and 1 output QDQ. Same as AveragePooling2DQuantizeWrapper.
    """

    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
        """
        Creates a wrapper to emulate quantization for the GlobalAveragePooling2D keras layer.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        super().__init__(layer, **kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, training=None):
        return super().call(inputs, training=training)


class MaxPooling2DQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
    """
    TRT Rule:
        Max pooling is precision-neutral. But unlike ReLU, input and output of max pooling will have different
            histograms which will lead to different calibration results.
        The recommendation is to let TensorRT optimize precision neutral ops.
        There are cases where adding Q/DQ before maxpool can enable additional optimization.

    Residual block v2: Add to MaxPool (branch1) and BN (branch2).

    Supports 1 input and 1 output QDQ. Same as AveragePooling2DQuantizeWrapper.
    """

    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
        """
        Creates a wrapper to emulate quantization for the MaxPooling2D keras layer.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        super().__init__(layer, **kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, training=None):
        return super().call(inputs, training=training)


class BatchNormalizationQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
    """
    TRT Rule:
        Keep batch normalization untouched, don't add Q/DQ to its input and not necessary to fold it before exporting
            graph. TensorRT supports Batch normalization folding. It can take a graph with batch normalization, fold it
            into previous convolution and create a new graph.
        If batch normalization is folded before exporting the graph, TensorRT can still import and execute the graph as
            it becomes regular convolutions.

    Exception for Residual block v2:
        BN-ReLU-Conv2D -> need to add Q/DQ before BN in order to run in INT8.
        In order to do that, we add a check in 'quantize_model()' to check if BN's parent is a Conv layer. If it is, set
            quantize_inputs to False. The reason why we don't add this check here is to allow the user to add QDQ nodes
            before BN if they so wish.

    Supports 1 input and 1 output QDQ. Same as AveragePooling2DQuantizeWrapper.
    """

    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
        """
        Creates a wrapper to emulate quantization for the BatchNormalization keras layer.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        super().__init__(layer, **kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, training=None):
        return super().call(inputs, training=training)


# ##############################################
# ########### Non-Weighted Layers ##############
# #### with Multiple Inputs, Single Output #####
# ##############################################


class NonWeightedBaseQuantizeWrapperForMultipleInputs(BaseQuantizeWrapper):
    """
    BaseQuantizeWrapper for non-weighted layers with Multiple Inputs: Concat, Add, and Multiply.
    Supports multiple inputs and 1 output QDQ. Similar to AveragePooling2D, except pooling supports only a single input.

    TRT Rule:
        Add Q/DQ to all inputs of the layer.
    """

    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
        """
        Creates a wrapper to emulate quantization for the keras layer.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        super().__init__(layer, **kwargs)

    def _should_quantization_this_index(self, i):
        if not self.quantize_specific_input_indices:
            return True
        else:
            # This is a small list so iterating makes sense
            for e in self.quantize_specific_input_indices:
                if e == i:
                    return True
                elif e >= self.num_inputs:
                    warnings.warn(
                        "{layer_name} has {num_inputs} inputs but quantization index {e} is passed.".format(
                            layer_name=self.layer.name, num_inputs=self.num_inputs, e=e
                        )
                    )
            return False

    def build(self, input_shape):
        super().build(input_shape)

        self.input_vars = []  # list of dictionaries
        self.num_inputs = len(input_shape)

        # By default input is per tensor quantization
        if self.quantize_inputs:
            # for concat input is list of Tensors
            layer_name_key_idx = 0
            for i in range(self.num_inputs):
                if self._should_quantization_this_index(i):
                    input_min_weight = self.layer.add_weight(
                        self.layer.name + "_ip{}_min".format(layer_name_key_idx),
                        shape=None,
                        initializer=tf.keras.initializers.Constant(-6.0),
                        trainable=False,
                    )
                    input_max_weight = self.layer.add_weight(
                        self.layer.name + "_ip{}_max".format(layer_name_key_idx),
                        shape=None,
                        initializer=tf.keras.initializers.Constant(6.0),
                        trainable=False,
                    )
                    self.input_vars.append(
                        {"min_var": input_min_weight, "max_var": input_max_weight}
                    )
                    layer_name_key_idx += 1

    def call(self, inputs, training=None):
        if training is None:
            training = tf.keras.backend.learning_phase()

        # Quantize inputs to the conv layer
        quantized_inputs = inputs[:]
        if self.quantize_inputs:
            input_vars_idx = 0
            for i in range(len(inputs)):
                if self._should_quantization_this_index(i):
                    quantized_inputs[i] = self._last_value_quantizer(
                        inputs[i],
                        training,
                        self.input_vars[input_vars_idx],
                        per_channel=False,
                    )
                    input_vars_idx += 1

        args = tf_inspect.getfullargspec(self.layer.call).args
        if "training" in args:
            outputs = self.layer.call(quantized_inputs, training=training)
        else:
            outputs = self.layer.call(quantized_inputs)

        return outputs


class MultiplyQuantizeWrapper(NonWeightedBaseQuantizeWrapperForMultipleInputs):
    """
    TRT Rule:
        Add Q/DQ to all inputs of Multiply layer in SE block.
    """

    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
        """
        Creates a wrapper to emulate quantization for the Multiply keras layer.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        super().__init__(layer, **kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, training=None):
        return super().call(inputs, training=training)


class ConcatenateQuantizeWrapper(NonWeightedBaseQuantizeWrapperForMultipleInputs):
    """
    TRT Rule:
        Add Q/DQ to all inputs.
        Alternative: If there is Q/DQ attached to the input of the op after concat, don't add Q/DQ to input to concat,
            let TensorRT pull from the next op.
    """

    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
        """
        Creates a wrapper to emulate quantization for the Concatenate keras layer.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        super().__init__(layer, **kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, training=None):
        return super().call(inputs, training=training)


class AddQuantizeWrapper(NonWeightedBaseQuantizeWrapperForMultipleInputs):
    """
    TRT Rule:
        If the add is NOT bias. Attach Q/DQ to all of its input.
        Exception: add in residual block. To trigger fusion, Attach Q/DQ to the residual being added to output of
            convolution.
    """

    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
        """
        Creates a wrapper to emulate quantization for the Add keras layer.
        Args:
          layer (tf.keras.layers.Layer): The keras layer to be quantized.
          **kwargs: Additional keyword arguments to be passed to the keras layer.
        """
        super().__init__(layer, **kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, training=None):
        return super().call(inputs, training=training)
