<commit_before># -*- coding: utf-8 -*-
"""Keras model parser.

@author: rbodo
"""
import os

import numpy as np
import tensorflow.keras.backend as k
from tensorflow.keras import models, metrics

from snntoolbox.parsing.utils import AbstractModelParser, has_weights, \
    fix_input_layer_shape, get_type, get_inbound_layers, get_outbound_layers, \
    get_custom_activations_dict, assemble_custom_dict, get_custom_layers_dict


class ModelParser(AbstractModelParser):

    def get_layer_iterable(self):
        return self.input_model.layers

    def get_type(self, layer):
        return get_type(layer)

    def get_batchnorm_parameters(self, layer):
        mean = k.get_value(layer.moving_mean)
        var = k.get_value(layer.moving_variance)
        var_eps_sqrt_inv = 1 / np.sqrt(var + layer.epsilon)
        gamma = np.ones_like(mean) if layer.gamma is None else \
            k.get_value(layer.gamma)
        beta = np.zeros_like(mean) if layer.beta is None else \
            k.get_value(layer.beta)
        axis = layer.axis
        if isinstance(axis, (list, tuple)):
            assert len(axis) == 1, "Multiple BatchNorm axes not understood."
            axis = axis[0]

        return [mean, var_eps_sqrt_inv, gamma, beta, axis]

    def get_inbound_layers(self, layer):
        return get_inbound_layers(layer)

    @property
    def layers_to_skip(self):
        # noinspection PyArgumentList
        return AbstractModelParser.layers_to_skip.fget(self)

    def has_weights(self, layer):
        return has_weights(layer)

    def initialize_attributes(self, layer=None):
        attributes = AbstractModelParser.initialize_attributes(self)
        attributes.update(layer.get_config())
        return attributes

    def get_input_shape(self):
        return \
            fix_input_layer_shape(self.get_layer_iterable()[0].input_shape)[1:]

    def get_output_shape(self, layer):
        return layer.output_shape

    def parse_sparse(self, layer, attributes):
        return self.parse_dense(layer, attributes)

    def parse_dense(self, layer, attributes):
        attributes['parameters'] = list(layer.get_weights())
        if layer.bias is None:
            attributes['parameters'].insert(
                1, np.zeros(layer.output_shape[1]))
            attributes['parameters'] = tuple(attributes['parameters'])
            attributes['use_bias'] = True

    def parse_sparse_convolution(self, layer, attributes):
        return self.parse_convolution(layer, attributes)

    def parse_convolution(self, layer, attributes):
        attributes['parameters'] = list(layer.get_weights())
        if layer.bias is None:
            attributes['parameters'].insert(1, np.zeros(layer.filters))
            attributes['parameters'] = tuple(attributes['parameters'])
            attributes['use_bias'] = True
        assert layer.data_format == k.image_data_format(), (
            "The input model was setup with image data format '{}', but your "
            "keras config file expects '{}'.".format(layer.data_format,
                                                     k.image_data_format()))

    def parse_sparse_depthwiseconvolution(self, layer, attributes):
        return self.parse_depthwiseconvolution(layer, attributes)

    def parse_depthwiseconvolution(self, layer, attributes):
        attributes['parameters'] = list(layer.get_weights())
        if layer.bias is None:
            a = 1 if layer.data_format == 'channels_first' else -1
            attributes['parameters'].insert(1, np.zeros(
                layer.depth_multiplier * layer.input_shape[a]))
            attributes['parameters'] = tuple(attributes['parameters'])
            attributes['use_bias'] = True

    def parse_pooling(self, layer, attributes):
        pass

    def get_activation(self, layer):

        return layer.activation.__name__

    def get_outbound_layers(self, layer):

        return get_outbound_layers(layer)

    def parse_concatenate(self, layer, attributes):
        pass

    @property
    def input_layer_name(self):
        return self.input_model.layers[0].name


def load(path, filename, **kwargs):
    """Load network from file.

    Parameters
    ----------

    path: str
        Path to directory where to load model from.

    filename: str
        Name of file to load model from.
<commit_msg>Fixed issue of naming the input layer of the parsed model correctly even if the original model has no dedicated input layer.<commit_after># -*- coding: utf-8 -*-
"""Keras model parser.

@author: rbodo
"""
import os

import numpy as np
import tensorflow.keras.backend as k
from tensorflow.keras import models, metrics

from snntoolbox.parsing.utils import AbstractModelParser, has_weights, \
    fix_input_layer_shape, get_type, get_inbound_layers, get_outbound_layers, \
    get_custom_activations_dict, assemble_custom_dict, get_custom_layers_dict


class ModelParser(AbstractModelParser):

    def get_layer_iterable(self):
        return self.input_model.layers

    def get_type(self, layer):
        return get_type(layer)

    def get_batchnorm_parameters(self, layer):
        mean = k.get_value(layer.moving_mean)
        var = k.get_value(layer.moving_variance)
        var_eps_sqrt_inv = 1 / np.sqrt(var + layer.epsilon)
        gamma = np.ones_like(mean) if layer.gamma is None else \
            k.get_value(layer.gamma)
        beta = np.zeros_like(mean) if layer.beta is None else \
            k.get_value(layer.beta)
        axis = layer.axis
        if isinstance(axis, (list, tuple)):
            assert len(axis) == 1, "Multiple BatchNorm axes not understood."
            axis = axis[0]

        return [mean, var_eps_sqrt_inv, gamma, beta, axis]

    def get_inbound_layers(self, layer):
        return get_inbound_layers(layer)

    @property
    def layers_to_skip(self):
        # noinspection PyArgumentList
        return AbstractModelParser.layers_to_skip.fget(self)

    def has_weights(self, layer):
        return has_weights(layer)

    def initialize_attributes(self, layer=None):
        attributes = AbstractModelParser.initialize_attributes(self)
        attributes.update(layer.get_config())
        return attributes

    def get_input_shape(self):
        return \
            fix_input_layer_shape(self.get_layer_iterable()[0].input_shape)[1:]

    def get_output_shape(self, layer):
        return layer.output_shape

    def parse_sparse(self, layer, attributes):
        return self.parse_dense(layer, attributes)

    def parse_dense(self, layer, attributes):
        attributes['parameters'] = list(layer.get_weights())
        if layer.bias is None:
            attributes['parameters'].insert(
                1, np.zeros(layer.output_shape[1]))
            attributes['parameters'] = tuple(attributes['parameters'])
            attributes['use_bias'] = True

    def parse_sparse_convolution(self, layer, attributes):
        return self.parse_convolution(layer, attributes)

    def parse_convolution(self, layer, attributes):
        attributes['parameters'] = list(layer.get_weights())
        if layer.bias is None:
            attributes['parameters'].insert(1, np.zeros(layer.filters))
            attributes['parameters'] = tuple(attributes['parameters'])
            attributes['use_bias'] = True
        assert layer.data_format == k.image_data_format(), (
            "The input model was setup with image data format '{}', but your "
            "keras config file expects '{}'.".format(layer.data_format,
                                                     k.image_data_format()))

    def parse_sparse_depthwiseconvolution(self, layer, attributes):
        return self.parse_depthwiseconvolution(layer, attributes)

    def parse_depthwiseconvolution(self, layer, attributes):
        attributes['parameters'] = list(layer.get_weights())
        if layer.bias is None:
            a = 1 if layer.data_format == 'channels_first' else -1
            attributes['parameters'].insert(1, np.zeros(
                layer.depth_multiplier * layer.input_shape[a]))
            attributes['parameters'] = tuple(attributes['parameters'])
            attributes['use_bias'] = True

    def parse_pooling(self, layer, attributes):
        pass

    def get_activation(self, layer):

        return layer.activation.__name__

    def get_outbound_layers(self, layer):

        return get_outbound_layers(layer)

    def parse_concatenate(self, layer, attributes):
        pass

    @property
    def input_layer_name(self):
        # Check if model has a dedicated input layer. If so, return its name.
        # Otherwise, the first layer might be a conv layer, so we return
        # 'input'.
        first_layer = self.input_model.layers[0]
        if 'Input' in self.get_type(first_layer):
            return first_layer.name
        else:
            return 'input'


def load(path, filename, **kwargs):
    """Load network from file.

    Parameters
    ----------

    path: str
        Path to directory where to load model from.

    filename: str
        Name of file to load model from.

{'fraction_non_alphanumeric': 0.09845814977973569, 'fraction_numerical': 0.0028634361233480176, 'mean_word_length': 2.793233082706767, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '818739', 'n_tokens_mistral': 2570, 'n_tokens_neox': 2484, 'n_words': 618}