#!/usr/bin/env python
# ******************************************************************************
# Copyright 2023 Brainchip Holdings Ltd.
#
# 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.
# ******************************************************************************
"""Functions to convert QuantizedConv2D to Akida.
"""
import numpy as np

from quantizeml.layers import (QuantizedConv2D, QuantizedReLU,
                               WeightQuantizer, AlignedWeightQuantizer)
from akida import InputConvolutional, Convolutional, LayerType

from ..akida_versions import AkidaVersion, get_akida_version
from .outputs import set_output_variables
from .padding import get_padding_info
from .layer_utils import parse_non_neural_layers


def _set_conv2d_variables(layer_ak, layer_k):
    """Computes and sets the variables for an Akida convolutional layer.

    This function converts the variables of a Keras layer and sets them into
    the corresponding variables of the equivalent Akida layer.

    Args:
        layer_ak (:obj:`akida.Layer`): the targeted akida layer.
        layer_k (:obj:`tf.keras.Layer`): the source quantized layer.
    """
    assert isinstance(layer_k, QuantizedConv2D)
    assert isinstance(layer_k.weight_quantizer, WeightQuantizer)
    if layer_k.use_bias:
        assert isinstance(layer_k.bias_quantizer, AlignedWeightQuantizer)

    # Get the weights
    weights = layer_k.weight_quantizer.qweights.value.fp.values.numpy()
    if layer_ak.parameters.layer_type == LayerType.Convolutional:
        # Flip W and H dimensions for conv. kernels (not input conv.)
        weights = np.flip(weights, axis=[0, 1])
    layer_ak.variables["weights"] = weights.astype(np.int8)

    # Set the bias (if there is one)
    if layer_k.use_bias:
        bias_quantizer = layer_k.bias_quantizer
        bias = bias_quantizer.qweights.value.values.numpy().astype(np.int32)
        bias_shift = bias_quantizer.shift.value.numpy().astype(np.uint8)
        target_version = get_akida_version()
        if target_version == AkidaVersion.v1:
            # Store bias into the threshold variable
            layer_ak.variables["threshold"] = -bias
        else:
            # Unshift the bias and store it
            layer_ak.variables["bias"] = bias >> bias_shift
            # Also store the bias shift
            layer_ak.variables["bias_shift"] = bias_shift


def _parse_convolutional(layer):
    """Parses a quantizeml.QuantizedConv2D parameters.

    Args:
        layer (:obj:`tf.keras.Layer`): the layer to parse.

    Returns:
        dict: the corresponding akida parameters.
    """
    assert isinstance(layer, QuantizedConv2D)

    # Padding value must be built in constructor
    padding, padding_value = get_padding_info(layer)

    # The only weight bitwidth supported is [4, 8]
    weight_bits = layer.weight_quantizer.bitwidth
    assert weight_bits in [4, 8]

    return dict(
        input_shape=tuple(int(x) for x in layer.input_shape[1:4]),
        padding=padding,
        padding_value=padding_value,
        kernel_size=(layer.kernel_size[0], layer.kernel_size[1]),
        filters=int(layer.kernel.shape[3]),
        weights_bits=layer.weight_quantizer.bitwidth,
        kernel_stride=(layer.strides[0], layer.strides[1]),
        activation=False,
        name=layer.name
    )


def convert_conv_block(model_ak, layers):
    """Converts an convolutional block into an akida convolutional layer.

    The expected sequence is:

    - QuantizedConv2D,
    - QuantizedMaxPooling2D (optional),
    - QuantizedRelU (optional).

    Args:
        model_ak (:obj:`akida.Model`): the target Akida model.
        layers (list(:obj:`tf.keras.Layer`)): the block layers.

    Return:
        int: the number of layers in the block or O if the first layer is not a QuantizedConv2D.
    """
    if len(layers) == 0 or not isinstance(layers[0], QuantizedConv2D):
        return 0

    if get_akida_version() == AkidaVersion.v2:
        return 0

    conv = layers[0]

    # Evaluate the convolutional layer parameters
    conv_params = _parse_convolutional(conv)
    next_layers = parse_non_neural_layers(layers, conv_params)

    if len(model_ak.layers) == 0 and conv_params["input_shape"][-1] in [1, 3]:
        # Only add InputConvolutional if nb_channels in [1, 3]
        conv_ak = InputConvolutional(**conv_params)
    elif len(model_ak.layers) != 0:
        # A Convolutional should be added if there is a previous layer (like InputData).
        # Unlike InputConvolutional, Convolutional does not receive the following parameters:
        for no_param in ["input_shape", "padding_value"]:
            conv_params.pop(no_param, None)
        conv_ak = Convolutional(**conv_params)
    else:
        # Convert loop will manage the other cases
        return 0

    # Add layer to the model to build its internal variables
    model_ak.add(conv_ak)
    # Set base variables
    _set_conv2d_variables(conv_ak, conv)
    for layer in next_layers:
        if isinstance(layer, QuantizedReLU):
            set_output_variables(conv_ak, layer.out_quantizer)

    return 1 + len(next_layers)
