#!/usr/bin/env python
# ******************************************************************************
# Copyright 2019 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.
# ******************************************************************************
"""
Parsing functions that are able to generate an Akida model from a keras model.
"""
import numpy as np
import tensorflow.keras.layers as layers
import yaml
from . import quantization_layers as qlayers


class ModelSkeleton:
    """
    Model skeleton builder from kernel layers.
    """
    def __init__(self, name="converted_model"):
        """
        :param name: name of generated akida model
        """
        self._model = {}
        self._model["Name"] = name
        self._layers = []

    def append(self, current, name):
        """
        Add a Keras layer to the mapping model.
        """
        if current is None or current == {}:
            return
        # If there was no activation params, set activations to none
        if ("thresholdFireBits" not in current and
                current["layerType"] != "inputData"):
            current["activations"] = "none"
        layer = {}
        layer["Name"] = name
        layer["Parameters"] = current.copy()
        self._layers.append(layer)
        current.clear()

    def yaml(self):
        """
        Convert the current object to yaml description string.
        """
        self._model["Layers"] = self._layers
        return yaml.dump(self._model)


def _parse_input_conv(layer, current):
    if not isinstance(layer, qlayers.QuantizedConv2D):
        raise TypeError(f"Layer {layer.name} was expected to be "
                        "QuantizedConv2D")
    current["layerType"] = "inputConvolutional"
    current["convolutionMode"] = layer.padding
    current["kernelHeight"] = layer.kernel_size[0]
    current["kernelWidth"] = layer.kernel_size[1]
    current["numNeurons"] = int(layer.kernel.shape[3])
    current["weightsBits"] = layer.quantizer.bitwidth
    current["inputChannels"] = int(layer.input_shape[3])
    current["strideX"] = layer.strides[1]
    current["strideY"] = layer.strides[0]


def _parse_conv(layer, current):
    dw_conv = qlayers.QuantizedDepthwiseConv2D
    if not isinstance(layer, qlayers.QuantizedConv2D) \
            or isinstance(layer, dw_conv):
        raise TypeError(f"Layer {layer.name} was expected to be "
                        "QuantizedConv2D")
    if layer.strides != (1,1):
        raise ValueError(f"Layer {layer.name} must have strides of (1,1). "
                        f"Received strides of {layer.strides}.")
    current["layerType"] = "convolutional"
    current["convolutionMode"] = layer.padding
    current["kernelHeight"] = layer.kernel_size[0]
    current["kernelWidth"] = layer.kernel_size[1]
    current["numNeurons"] = int(layer.kernel.shape[3])
    current["weightsBits"] = layer.quantizer.bitwidth


def _parse_separable_conv(layer, current):
    if not isinstance(layer, qlayers.QuantizedSeparableConv2D):
        raise TypeError(f"Layer {layer.name} was expected to be "
                        "QuantizedSeparableConv2D")
    current["layerType"] = "separableConvolutional"
    current["convolutionMode"] = layer.padding
    current["kernelHeight"] = layer.kernel_size[0]
    current["kernelWidth"] = layer.kernel_size[1]
    # num neurons is set to the number of filters of the depthwise
    current["numNeurons"] = int(layer.depthwise_kernel.shape[2])
    # For now depthwise multiplier is not supported
    if layer.depthwise_kernel.shape[3] != 1:
        raise RuntimeError(f"Layer {layer.name}: unsupported depthwise "
                           "multiplier")
    current["numPointwiseNeurons"] = int(layer.pointwise_kernel.shape[3])
    current["weightsBits"] = layer.quantizer.bitwidth


def _parse_dense(layer, current):
    if not isinstance(layer, qlayers.QuantizedDense):
        raise TypeError(f"Layer {layer.name} was expected to be "
                        "QuantizedDense")
    current["layerType"] = "fullyConnected"
    current["numNeurons"] = layer.units
    current["weightsBits"] = layer.quantizer.bitwidth


def _parse_max_pooling(layer, current):
    if not isinstance(layer, layers.MaxPooling2D):
        raise TypeError(f"Layer {layer.name} was expected to be MaxPooling2D")
    current["poolingType"] = "max"
    current["poolingHeight"] = layer.pool_size[0]
    current["poolingWidth"] = layer.pool_size[1]
    current["poolStrideY"] = layer.strides[0]
    current["poolStrideX"] = layer.strides[1]
    # Check the padding is identical to the one in the convolution
    if current["convolutionMode"] != layer.padding:
        raise ValueError(f"Layer {layer.name} has a padding {layer.padding},"
                          "which is different from the padding of the related"
                          "convolutional layer")

def _parse_global_average_pooling(layer, current):
    if not isinstance(layer, layers.GlobalAveragePooling2D):
        raise TypeError(f"Layer {layer.name} was expected to be "
                        "GlobalAveragePooling2D")
    current["poolingType"] = "average"
    current["poolingSize"] = -1


def generate_model(model, input_scaling, input_is_sparse=False):
    """
    Create an akida model that maps an existing keras model

    :param model: a tf.keras model
    :type model: tf.keras.Model
    :param input_scaling: value of the input scaling, alpha and beta
    :type input_scaling: 2 elements tuple
    :input_is_sparse: if True, input will be an inputData layer,
                      otherwise it will be inputConvolutional
    :input_is_sparse: bool

    The relationship between Keras and Akida inputs is:

        input_akida = alpha * input_keras + beta

    :return: a string with the model in yaml
    """
    for layer in model.layers:
        if hasattr(layer, 'data_format'):
            if layer.data_format == "channels_first":
                raise RuntimeError("unsupported data format channels_first")

    model_skeleton = ModelSkeleton(model.name)

    # Current layer config
    current = {}

    # First we need to map the input layer
    start_layer = 0
    # If first layer is input layer, skip it
    if isinstance(model.layers[0], layers.InputLayer):
        start_layer = 1
    input_layer = model.layers[start_layer]
    current_name = input_layer.name
    beta = input_scaling[1]
    if input_is_sparse:
        if beta != 0:
            raise ValueError(
                "beta should be 0 when input_is_sparse is True.")
        current["layerType"] = "inputData"
        current["inputFeatures"] = int(input_layer.input_shape[3])
        current_name += "_input"
        # We do not set these attributes: packetSize, accumulate.
    else:
        _parse_input_conv(input_layer, current)
        current["paddingValue"] = int(beta)
        start_layer += 1
    current["inputHeight"] = int(input_layer.input_shape[1])
    current["inputWidth"] = int(input_layer.input_shape[2])

    # Loop on layers
    for i in range(start_layer, len(model.layers)):
        layer = model.layers[i]
        if isinstance(layer, qlayers.QuantizedConv2D):
            model_skeleton.append(current, current_name)
            _parse_conv(layer, current)
            current_name = layer.name
            continue
        if isinstance(layer, qlayers.QuantizedSeparableConv2D):
            model_skeleton.append(current, current_name)
            _parse_separable_conv(layer, current)
            current_name = layer.name
            continue
        if isinstance(layer, qlayers.QuantizedDense):
            model_skeleton.append(current, current_name)
            _parse_dense(layer, current)
            current_name = layer.name
            continue
        # Parse pooling params/layer
        layers_with_pooling = ["separableConvolutional",
                               "convolutional",
                               "inputConvolutional"]
        if isinstance(layer, layers.MaxPooling2D):
            if current["layerType"] not in layers_with_pooling:
                model_skeleton.append(current, current_name)
                current["layerType"] = "pooling"
                current_name = layer.name
            _parse_max_pooling(layer, current)
            continue
        if isinstance(layer, layers.GlobalAveragePooling2D):
            if current["layerType"] not in layers_with_pooling:
                model_skeleton.append(current, current_name)
                current["layerType"] = "pooling"
                current_name = layer.name
            _parse_global_average_pooling(layer, current)
            continue
        if isinstance(layer, qlayers.BaseQuantizedActivation):
            current["thresholdFireBits"] = layer.bitwidth
        # Allow flatten before a dense layer
        if isinstance(layer, layers.Flatten):
            try:
                if isinstance(model.layers[i+1], qlayers.QuantizedDense):
                    continue
            except IndexError:
                pass
            raise RuntimeError("Flatten layer only supported before a Dense "
                               "one")
        # Allow some other layers useful in keras but that will be discarded
        # or ignored during conversion
        ignore_list = (layers.Activation,
                       layers.Reshape,
                       layers.Dropout,
                       layers.Softmax,
                       layers.BatchNormalization)
        if isinstance(layer, ignore_list):
            continue

        # If you got here it means the layer is not recognised: raise an error.
        raise RuntimeError(f"layer {layer.name}: unsupported type "
                           f"{layer.__class__.__name__}.")

    # Append last parsed layer if any
    model_skeleton.append(current, current_name)

    return model_skeleton.yaml()
