# Path: snntoolbox/simulation/utils.py
# def get_layer_synaptic_operations(spiketrains_b_l, fanout):
#     """
#     Return total number of synaptic operations in the layer for a batch of
#     samples.
# 
#     Parameters
#     ----------
# 
#     spiketrains_b_l: ndarray
#         Batch of spiketrains of a layer. Shape: (batch_size, layer_shape)
#     fanout: Union[int, ndarray]
#         Number of outgoing connections per neuron. Can be a single integer, or
#         an array of the same shape as the layer, if the fanout varies from
#         neuron to neuron (as is the case in convolution layers with
#         stride > 1).
# 
#     Returns
#     -------
# 
#     layer_ops: int
#         The total number of operations in the layer for a batch of samples.
#     """
# 
#     if np.isscalar(fanout):
#         return np.array([np.count_nonzero(s) for s in spiketrains_b_l]) * \
#             fanout
#     elif hasattr(fanout, 'shape'):  # For conv layers with stride > 1
#         return np.array([np.sum(fanout[s != 0]) for s in spiketrains_b_l])
#     else:
#         raise TypeError("The 'fanout' parameter should either be integer or "
#                         "ndarray.")
# 
# def remove_name_counter(name_in):
#     """
#     Tensorflow adds a counter to layer names, e.g. <name>/kernel:0 ->
#     <name>_0/kernel:0. Need to remove this _0.
#     The situation gets complicated because SNN toolbox assigns layer names
#     that contain the layer shape, e.g. 00Conv2D_3x32x32. In addition,
#     we may get another underscore in the parameter name, e.g.
#     00DepthwiseConv2D_3X32x32_0/depthwise_kernel:0.
#     """
# 
#     if '_' not in name_in or '/' not in name_in:
#         return name_in
# 
#     split_dash = str(name_in).split('/')
#     assert len(split_dash) == 2, "Variable name must not contain '/'."
#     # We are only interested in the part before the /.
#     split_underscore = split_dash[0].split('_')
#     # The first '_' is assigned by SNN toolbox and should be kept.
#     return (split_underscore[0] + '_' + split_underscore[1] + '/' +
#             split_dash[1])


import tensorflow as tf
import numpy as np
from snntoolbox.simulation.target_simulators.\
    INI_temporal_mean_rate_target_sim import SNN as SNN_
from snntoolbox.simulation.utils import get_layer_synaptic_operations, \
    remove_name_counter
        from snntoolbox.utils.utils import echo

    Aims at simulating the network on a self-implemented Integrate-and-Fire
    simulator using a timestepped approach.

    Attributes
    ----------

    snn: keras.models.Model
        Keras model. This is the output format of the compiled spiking model
        because INI simulator runs networks of layers that are derived from
        Keras layer base classes.
    """

    def __init__(self, config, queue=None):

        SNN_.__init__(self, config, queue)

        self.num_bits = self.config.getint('conversion', 'num_bits')

    def compile(self):

        self.snn = tf.keras.models.Model(
            self._input_images,
            self._spiking_layers[self.parsed_model.layers[-1].name])
        self.snn.compile('sgd', 'categorical_crossentropy', ['accuracy'])

        # Tensorflow 2 lists all variables as weights, including our state
        # variables (membrane potential etc). So a simple
        # snn.set_weights(parsed_model.get_weights()) does not work any more.
        # Need to extract the actual weights here:

# Path: snntoolbox/simulation/utils.py
# def get_layer_synaptic_operations(spiketrains_b_l, fanout):
#     """
#     Return total number of synaptic operations in the layer for a batch of
#     samples.
# 
#     Parameters
#     ----------
# 
#     spiketrains_b_l: ndarray
#         Batch of spiketrains of a layer. Shape: (batch_size, layer_shape)
#     fanout: Union[int, ndarray]
#         Number of outgoing connections per neuron. Can be a single integer, or
#         an array of the same shape as the layer, if the fanout varies from
#         neuron to neuron (as is the case in convolution layers with
#         stride > 1).
# 
#     Returns
#     -------
# 
#     layer_ops: int
#         The total number of operations in the layer for a batch of samples.
#     """
# 
#     if np.isscalar(fanout):
#         return np.array([np.count_nonzero(s) for s in spiketrains_b_l]) * \
#             fanout
#     elif hasattr(fanout, 'shape'):  # For conv layers with stride > 1
#         return np.array([np.sum(fanout[s != 0]) for s in spiketrains_b_l])
#     else:
#         raise TypeError("The 'fanout' parameter should either be integer or "
#                         "ndarray.")
# 
# def remove_name_counter(name_in):
#     """
#     Tensorflow adds a counter to layer names, e.g. <name>/kernel:0 ->
#     <name>_0/kernel:0. Need to remove this _0.
#     The situation gets complicated because SNN toolbox assigns layer names
#     that contain the layer shape, e.g. 00Conv2D_3x32x32. In addition,
#     we may get another underscore in the parameter name, e.g.
#     00DepthwiseConv2D_3X32x32_0/depthwise_kernel:0.
#     """
# 
#     if '_' not in name_in or '/' not in name_in:
#         return name_in
# 
#     split_dash = str(name_in).split('/')
#     assert len(split_dash) == 2, "Variable name must not contain '/'."
#     # We are only interested in the part before the /.
#     split_underscore = split_dash[0].split('_')
#     # The first '_' is assigned by SNN toolbox and should be kept.
#     return (split_underscore[0] + '_' + split_underscore[1] + '/' +
#             split_dash[1])


# Path: snntoolbox/simulation/target_simulators/INI_temporal_pattern_target_sim.py
import tensorflow as tf
import numpy as np
from snntoolbox.simulation.target_simulators.\
    INI_temporal_mean_rate_target_sim import SNN as SNN_
from snntoolbox.simulation.utils import get_layer_synaptic_operations, \
    remove_name_counter
        from snntoolbox.utils.utils import echo


    Aims at simulating the network on a self-implemented Integrate-and-Fire
    simulator using a timestepped approach.

    Attributes
    ----------

    snn: keras.models.Model
        Keras model. This is the output format of the compiled spiking model
        because INI simulator runs networks of layers that are derived from
        Keras layer base classes.
    """

    def __init__(self, config, queue=None):

        SNN_.__init__(self, config, queue)

        self.num_bits = self.config.getint('conversion', 'num_bits')

    def compile(self):

        self.snn = tf.keras.models.Model(
            self._input_images,
            self._spiking_layers[self.parsed_model.layers[-1].name])
        self.snn.compile('sgd', 'categorical_crossentropy', ['accuracy'])

        # Tensorflow 2 lists all variables as weights, including our state
        # variables (membrane potential etc). So a simple
        # snn.set_weights(parsed_model.get_weights()) does not work any more.
        # Need to extract the actual weights here:

        parameter_map = {remove_name_counter(p.name): p for p in