#!/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.
# ******************************************************************************
"""
Tools to unfuse SeparableConv2D layers from a model.
"""

__all__ = ["unfuse_sepconv2d"]

from copy import deepcopy

from keras.models import Sequential
from keras.layers import SeparableConv2D
from quantizeml.models.utils import apply_weights_to_model
from quantizeml.models.transforms.transforms_utils import (
    get_layer_index, get_layers_by_type, update_inbound)


def _find_sepconv_layers(model):
    """ Retrieves SeparableConv2D layers.

    Args:
        model (keras.Model): a model

    Returns:
        dict: map between a SeparableConv2D and the layer that follows
    """
    map_sepconv_next = {}

    # Get all SeparableConv2D layers present in the model
    sep_conv_layers = get_layers_by_type(model, SeparableConv2D)

    for sep_conv in sep_conv_layers:
        # Limit support to single inbound/outbound
        outbounds = sep_conv.outbound_nodes
        if len(sep_conv.inbound_nodes) != 1 or len(outbounds) != 1:
            continue

        following_layer = outbounds[0].layer
        # At this point the SeparableConv2D is a valid candidate
        map_sepconv_next[sep_conv] = following_layer

    # Check if the last model layer is a SeparableConv2D if it's the case add it
    # to dict but with None as a special value (to indicate that it has no outbound layer)
    last_layer = model.layers[-1]
    if isinstance(last_layer, SeparableConv2D):
        map_sepconv_next[last_layer] = None
    return map_sepconv_next


def _get_depthwise_conf(sep_layer):
    """ Helper function to unfuse a SeparableConv2D layer. This function returns the
        corresponding DepthwiseConv2D layer configuration and its generated name.

    Args:
        sep_layer (keras.Layer): a keras SeparableConv2D layer

    Returns:
        dict: the DepthwiseConv2D layer configuration
    """
    assert isinstance(sep_layer, SeparableConv2D), ("_get_depthwise_conf accepts only "
                                                    "SeparableConv2D layers. "
                                                    f"Got {type(sep_layer)}")
    dw_name = "dw_" + sep_layer.name
    sep_layer_config = sep_layer.get_config()
    depthwise_initializer = sep_layer_config.get("depthwise_initializer",
                                                 {'class_name': 'GlorotUniform',
                                                  'config': {'seed': None}})
    depthwise_regularizer = sep_layer_config.get("depthwise_regularizer", None)
    depthwise_constraint = sep_layer_config.get("depthwise_constraint", None)

    dw_layer_conf = {'class_name': 'DepthwiseConv2D',
                     'config': {'name': dw_name, 'trainable': sep_layer.trainable,
                                'dtype': sep_layer.dtype,
                                'kernel_size': sep_layer.kernel_size,
                                'strides': sep_layer.strides,
                                'padding': sep_layer.padding, 'data_format': sep_layer.data_format,
                                'dilation_rate': sep_layer.dilation_rate,
                                'groups': sep_layer.groups, 'activation': 'linear',
                                'use_bias': False,
                                'bias_initializer': {'class_name': 'Zeros', 'config': {}},
                                'bias_regularizer': None,
                                'bias_constraint': None,
                                'activity_regularizer': None,
                                'depth_multiplier': sep_layer.depth_multiplier,
                                'depthwise_initializer': depthwise_initializer,
                                'depthwise_regularizer': depthwise_regularizer,
                                'depthwise_constraint': depthwise_constraint}}

    return dw_layer_conf


def _get_pointwise_conf(sep_layer):
    """ Helper function to unfuse a SeparableConv2D layer. This function returns the
        corresponding (pointwise)Conv2D layer configuration and its generated name.

    Args:
        sep_layer (keras.Layer): a keras SeparableConv2D layer

    Returns:
        dict: the (pointwise)Conv2D layer configuration
    """
    assert isinstance(sep_layer, SeparableConv2D), ("_get_pointwise_conf accepts only "
                                                    "SeparableConv2D layers. "
                                                    f"Got {type(sep_layer)}")
    pw_name = "pw_" + sep_layer.name
    sep_layer_config = sep_layer.get_config()
    pointwise_initializer = sep_layer_config.get("pointwise_initializer",
                                                 {'class_name': 'GlorotUniform',
                                                  'config': {'seed': None}})
    pointwise_regularizer = sep_layer_config.get("pointwise_regularizer", None)
    pointwise_constraint = sep_layer_config.get("pointwise_constraint", None)
    bias_initializer = sep_layer_config.get("bias_initializer",
                                            {'class_name': 'Zeros', 'config': {}})
    bias_regularizer = sep_layer_config.get("bias_regularizer", None)
    bias_constraint = sep_layer_config.get("bias_constraint", None)
    activation = sep_layer_config.get("activation", "linear")
    activity_regularizer = sep_layer_config.get("activity_regularizer", None)

    pw_layer_conf = {'class_name': 'Conv2D',
                     'config': {'name': pw_name, 'trainable': sep_layer.trainable,
                                'dtype': sep_layer.dtype, 'filters': sep_layer.filters,
                                'kernel_size': (1, 1), 'strides': (1, 1),
                                'padding': 'valid', 'data_format': sep_layer.data_format,
                                'dilation_rate': (1, 1), 'groups': sep_layer.groups,
                                'activation': activation, 'use_bias': sep_layer.use_bias,
                                'kernel_initializer': pointwise_initializer,
                                'bias_initializer': bias_initializer,
                                'kernel_regularizer': pointwise_regularizer,
                                'bias_regularizer': bias_regularizer,
                                'activity_regularizer': activity_regularizer,
                                'kernel_constraint': pointwise_constraint,
                                'bias_constraint': bias_constraint}}

    return pw_layer_conf


def _get_unfused_sepconv_model(model, map_sepconv_layers):
    """Edits the model configuration to unfuse SeparableConv2D layers and rebuilds a model.
    Returns also a dict mapping the new unfused layers variable names with their
    corresponding values (i.e from the main SeparableConv2D layer).

    Args:
        model (keras.Model): a model
        map_sepconv_layers (dict): map between a SeparableConv2D and the layer that follows (None
            if the SeparableConv2D layer is the last model layer)

    Returns:
        tuple (keras.Model, dict): (an updated model with unfused SeparableConv2D layers,
        new layers variables)
    """
    # get_config documentation mentions that a copy should be made when planning to modify the
    # config
    config = deepcopy(model.get_config())
    layers = config['layers']

    # Create an empty dict that will be populated with the new layers variables names
    # and their corresponding values
    sep_layers_vars_dict = {}

    for sep_conv, next_layer in map_sepconv_layers.items():
        # Get the index of the SeparableConv2D layer
        sep_conv_index = get_layer_index(layers, sep_conv.name)

        # Get the original SeparableConv2D layer inbound informations if available
        sep_conv_inbounds_conf = layers[sep_conv_index].get('inbound_nodes', None)

        # Get DepthwiseConv2D conf from the main SeparableConv2D layer
        dw_conf = _get_depthwise_conf(sep_conv)
        # Replace the SeparableConv2D by a DepthwiseConv2D
        layers[sep_conv_index] = dw_conf

        # And insert after it a Conv2D layer (i.e pointwise conv2d)
        pw_conf = _get_pointwise_conf(sep_conv)
        layers.insert(sep_conv_index + 1, pw_conf)

        # Get new layers names
        dw_name = dw_conf['config']['name']
        pw_name = pw_conf['config']['name']
        # Get new layers var names
        dw_kernel_name = dw_name+"/depthwise_kernel:0"
        pw_kernel_name = pw_name+"/kernel:0"
        pw_bias_name = pw_name+"/bias:0"
        # Map new vars with the corresponding ones from the SeparableConv2D layer
        sep_layers_vars_dict[dw_kernel_name] = sep_conv.depthwise_kernel
        sep_layers_vars_dict[pw_kernel_name] = sep_conv.pointwise_kernel
        sep_layers_vars_dict[pw_bias_name] = sep_conv.bias

        # For sequential model, the changes stop here: the SeparableConv2D layers will simply be
        # replaced by an equivalent DepthwiseConv2D + (pointwise)Conv2D layers. For other models,
        # the layers inbounds/outbounds must be rebuilt.
        if isinstance(model, Sequential):
            continue

        # Retrieve the new DepthwiseConv2D index
        dw_layer_index = get_layer_index(layers, dw_name)
        layers[dw_layer_index]['name'] = dw_name
        # Set the new DepthwiseConv2D layer inbounds (those are the main SeparableConv2D inbounds)
        layers[dw_layer_index]['inbound_nodes'] = sep_conv_inbounds_conf

        # Retrieve the new (pointwise)Conv2D index (dw_layer_index + 1)
        pw_layer_index = dw_layer_index + 1
        layers[pw_layer_index]['name'] = pw_name
        # Set the new (pointwise)Conv2D layer inbounds (its the new DepthwiseConv2D layer)
        # tfmot code: 'inbound_nodes' is a nested list where first element is the inbound layername,
        # e.g: [[['conv1', 0, 0, {} ]]]
        layers[pw_layer_index]['inbound_nodes'] = [[[dw_name, 0, 0, {}]]]

        if next_layer:
            # Retrieve the next layer index
            next_layer_index = get_layer_index(layers, next_layer.name)
            # Update the next layer inbound layer (i.e the main SeparableConv2D layer) by the
            # new (pointwise)Conv2D
            update_inbound(layers[next_layer_index], sep_conv.name, pw_name)
        else:
            config['output_layers'] = [[pw_name, 0, 0]]

    # Reconstruct model from the config, using the cloned layers
    return model.from_config(config), sep_layers_vars_dict


def unfuse_sepconv2d(model):
    """ Unfuse the SeparableConv2D layers of a model by replacing them with an equivalent
    DepthwiseConv2D + (pointwise)Conv2D layers.

    Args:
        model (keras.Model): the model to update

    Returns:
        keras.Model: the original model or a new model with unfused SeparableConv2D layers
    """
    # Find SeparableConv2D
    map_sepconv_layers = _find_sepconv_layers(model)

    # When there are no valid candidates, return the original model
    if not map_sepconv_layers:
        return model

    # Rebuild a model with unfused SeparableConv2D by editing the configuration
    updated_model, sep_layers_vars_dict = _get_unfused_sepconv_model(model, map_sepconv_layers)

    # Load original weights
    variables_dict = {}

    for layer in model.layers:
        if not isinstance(layer, SeparableConv2D):
            for var in layer.variables:
                variables_dict[var.name] = var

    variables_dict.update(sep_layers_vars_dict)
    apply_weights_to_model(updated_model, variables_dict, False)

    return updated_model
