# Copyright (c) 2018 Uber Technologies, Inc.
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Layer, InputSpec
import tensorflow.keras.backend as K
from tensorflow.python.keras.utils import conv_utils
from src.subspace.keras_ext.rproj_layers_util import _convert_string_dtype
# from keras.backend.tensorflow_backend import _convert_string_dtype
from tensorflow.keras import regularizers, constraints, initializers, activations

###########
#
# Low Rank Basis Layers
#
# These layers are modified versions of standard Keras layers that
# accept an OffsetCreator*Proj to create offsets from a weight basis
# in a Dense/Sparse/Fastfood agnostic manner.
#
###########

class LowRankBasisLayer(Layer):
    '''Smarter version of Layer...'''

    def __init__(self, offset_creator_class, weight_basis, *args, **kwargs):
        super(LowRankBasisLayer, self).__init__(*args, **kwargs)
        # offset_creator is an object that creates theta offsets
        self.offset_creator = offset_creator_class()
        self.weight_basis = weight_basis

        # These may or may not be used by subclasses
        #self._basis_matrices = []
        #self._basis_matrix_normalizers = []

    # TODO check for use of basis_matrices

    @property
    def basis_matrices(self):
        print('USED HERE basis_matrices')
        return self._basis_matrices

    # TODO check for use of basis_matrix_normalizers
    @property
    def basis_matrix_normalizers(self):
        print('USED HERE basis_matrix_normalizers')
        return self._basis_matrix_normalizers
    
    def add_weight(self,
                   name,
                   shape,
                   dtype=None,
                   initializer=None,
                   regularizer=None,
                   trainable=True,
                   constraint=None):
        '''Version of add_weight that creates a weight theta by instantiating
        theta_0 and then adding to it an offset from the member
        offset_creator.
        '''
        initializer = initializers.get(initializer)
        if dtype is None:
            dtype = K.floatx()

        # Create Theta_0
        value_0 = initializer(shape)
        theta_0 = tf.Variable(value_0, trainable=False, dtype=_convert_string_dtype(dtype), name='%s_theta0' % name)
        if isinstance(value_0, np.ndarray):
            theta_0._keras_shape = value_0.shape
        elif hasattr(value_0, 'get_shape'):
            theta_0._keras_shape = tuple(map(int, value_0.get_shape()))
        theta_0._uses_learning_phase = False

        # Call offset creator
        exec = self.offset_creator.create_theta_offset(self.weight_basis,
                                                                                      theta_0.get_shape(),
                                                                                      dtype=dtype,
                                                                                      name=name)
        non_trainable_weights = exec.ww


        # if regularizer is not None:
        #     self.add_loss(regularizer(theta))
        # if constraint is not None:
        #     self.constraints[theta] = constraint
        #self._base_thetas.append(theta_0)
        #self._basis_matrices.append(ww)
        #self._non_trainable_weights.extend([theta_0, ww])
        self._non_trainable_weights.extend([theta_0] + [non_trainable_weights])
        return theta_0, exec

    def add_non_trainable_weight(self,
                                 name,
                                 shape,
                                 dtype=None,
                                 initializer=None,
                                 regularizer=None,
                                 constraint=None):
        '''Adds a weight variable to the layer.
        # Arguments
            name: String, the name for the weight variable.
            shape: The shape tuple of the weight.
            dtype: The dtype of the weight.
            initializer: An Initializer instance (callable).
            regularizer: An optional Regularizer instance.
            trainable: A boolean, whether the weight should
                be trained via backprop or not (assuming
                that the layer itself is also trainable).
            constraint: An optional Constraint instance.
        # Returns
            The created weight variable.
        '''
        initializer = initializers.get(initializer)
        if dtype is None:
            dtype = K.floatx()
        weight = K.variable(initializer(shape), dtype=dtype, name=name)
        if regularizer is not None:
            self.add_loss(regularizer(weight))
        if constraint is not None:
            self.constraints[weight] = constraint
        self._non_trainable_weights.append(weight)
        return weight


class RProjDense(LowRankBasisLayer):
    '''RProj version of Dense.'''

    def __init__(self, offset_creator_class, weight_basis,
                 units,
                 activation=None,
                 use_bias=True,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 kernel_regularizer=None,
                 bias_regularizer=None,
                 activity_regularizer=None,
                 kernel_constraint=None,
                 bias_constraint=None,
                 **kwargs):
        if 'input_shape' not in kwargs and 'input_dim' in kwargs:
            kwargs['input_shape'] = (kwargs.pop('input_dim'),)
        super(RProjDense, self).__init__(offset_creator_class, weight_basis, **kwargs)
        self.units = units
        self.activation = activations.get(activation)
        self.use_bias = use_bias
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.bias_initializer = initializers.get(bias_initializer)
        self.kernel_regularizer = regularizers.get(kernel_regularizer)
        self.bias_regularizer = regularizers.get(bias_regularizer)
        self.activity_regularizer = regularizers.get(activity_regularizer)
        self.kernel_constraint = constraints.get(kernel_constraint)
        self.bias_constraint = constraints.get(bias_constraint)
        self.input_spec = InputSpec(min_ndim=2)
        self.supports_masking = True

    def build(self, input_shape):
        assert len(input_shape) >= 2
        input_dim = input_shape[-1]

        self.kernel = self.add_weight(shape=(input_dim, self.units),
                                      initializer=self.kernel_initializer,
                                      name='kernel',
                                      regularizer=self.kernel_regularizer,
                                      constraint=self.kernel_constraint)
        if self.use_bias:
            self.bias = self.add_weight(shape=(self.units,),
                                        initializer=self.bias_initializer,
                                        name='bias',
                                        regularizer=self.bias_regularizer,
                                        constraint=self.bias_constraint)
        else:
            self.bias = None
        self.input_spec = InputSpec(min_ndim=2, axes={-1: input_dim})
        self.built = True

    def call(self, inputs, **kwargs):

        kt0, eproj = self.kernel
        k = tf.add(kt0, eproj())

        bt0, eprojb = self.bias
        b = tf.add(bt0, eprojb())

        # Normal dense functionality
        output = K.dot(inputs, k)
        if self.use_bias:
            output = K.bias_add(output, b)
        if self.activation is not None:
            output = self.activation(output)
        return output

    def compute_output_shape(self, input_shape):
        assert input_shape and len(input_shape) >= 2
        assert input_shape[-1]
        output_shape = list(input_shape)
        output_shape[-1] = self.units
        return tuple(output_shape)



class _RProjConv(LowRankBasisLayer):
    '''Abstract nD convolution layer (private, used as implementation base).

    Only the intrinsic parameters (RProj) are Trainable.'''

    def __init__(self, offset_creator_class, weight_basis,
                 rank,
                 filters,
                 kernel_size,
                 strides=1,
                 padding='valid',
                 data_format=None,
                 dilation_rate=1,
                 activation=None,
                 use_bias=True,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 kernel_regularizer=None,
                 bias_regularizer=None,
                 activity_regularizer=None,
                 kernel_constraint=None,
                 bias_constraint=None,
                 **kwargs):
        super(_RProjConv, self).__init__(offset_creator_class, weight_basis, **kwargs)
        self.rank = rank
        self.filters = filters
        self.kernel_size = conv_utils.normalize_tuple(kernel_size, rank, 'kernel_size')
        self.strides = conv_utils.normalize_tuple(strides, rank, 'strides')
        self.padding = conv_utils.normalize_padding(padding)
        self.data_format = conv_utils.normalize_data_format(data_format)
        self.dilation_rate = conv_utils.normalize_tuple(dilation_rate, rank, 'dilation_rate')
        self.activation = activations.get(activation)
        self.use_bias = use_bias
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.bias_initializer = initializers.get(bias_initializer)
        self.kernel_regularizer = regularizers.get(kernel_regularizer)
        self.bias_regularizer = regularizers.get(bias_regularizer)
        self.activity_regularizer = regularizers.get(activity_regularizer)
        self.kernel_constraint = constraints.get(kernel_constraint)
        self.bias_constraint = constraints.get(bias_constraint)
        self.input_spec = InputSpec(ndim=self.rank + 2)

    def build(self, input_shape):
        if self.data_format == 'channels_first':
            channel_axis = 1
        else:
            channel_axis = -1
        if input_shape[channel_axis] is None:
            raise ValueError('The channel dimension of the inputs '
                             'should be defined. Found `None`.')
        input_dim = input_shape[channel_axis]
        kernel_shape = self.kernel_size + (input_dim, self.filters)

        self.kernel = self.add_weight(shape=kernel_shape,
                                      initializer=self.kernel_initializer,
                                      name='kernel',
                                      regularizer=self.kernel_regularizer,
                                      constraint=self.kernel_constraint)
        if self.use_bias:
            self.bias = self.add_weight(shape=(self.filters,),
                                        initializer=self.bias_initializer,
                                        name='bias',
                                        regularizer=self.bias_regularizer,
                                        constraint=self.bias_constraint)
        else:
            self.bias = None
        # Set input spec.
        self.input_spec = InputSpec(ndim=self.rank + 2,
                                    axes={channel_axis: input_dim})
        self.built = True

    def call(self, inputs):
        if self.rank == 1:
            outputs = K.conv1d(
                inputs,
                self.kernel,
                strides=self.strides[0],
                padding=self.padding,
                data_format=self.data_format,
                dilation_rate=self.dilation_rate[0])
        if self.rank == 2:
            outputs = K.conv2d(
                inputs,
                self.kernel,
                strides=self.strides,
                padding=self.padding,
                data_format=self.data_format,
                dilation_rate=self.dilation_rate)
        if self.rank == 3:
            outputs = K.conv3d(
                inputs,
                self.kernel,
                strides=self.strides,
                padding=self.padding,
                data_format=self.data_format,
                dilation_rate=self.dilation_rate)

        if self.use_bias:
            outputs = K.bias_add(
                outputs,
                self.bias,
                data_format=self.data_format)

        if self.activation is not None:
            return self.activation(outputs)
        return outputs

    def compute_output_shape(self, input_shape):
        if self.data_format == 'channels_last':
            space = input_shape[1:-1]
            new_space = []
            for i in range(len(space)):
                new_dim = conv_utils.conv_output_length(
                    space[i],
                    self.kernel_size[i],
                    padding=self.padding,
                    stride=self.strides[i],
                    dilation=self.dilation_rate[i])
                new_space.append(new_dim)
            return (input_shape[0],) + tuple(new_space) + (self.filters,)
        if self.data_format == 'channels_first':
            space = input_shape[2:]
            new_space = []
            for i in range(len(space)):
                new_dim = conv_utils.conv_output_length(
                    space[i],
                    self.kernel_size[i],
                    padding=self.padding,
                    stride=self.strides[i],
                    dilation=self.dilation_rate[i])
                new_space.append(new_dim)
            return (input_shape[0], self.filters) + tuple(new_space)


class RProjConv2D(_RProjConv):
    '''Low Rank Basis Conv2D
    Filters if number of filters, output dimension is filters
    TODO: Documentation / unit tests
    '''

    def __init__(self, offset_creator_class, weight_basis,
                 filters,
                 kernel_size,
                 strides=(1, 1),
                 padding='valid',
                 data_format=None,
                 dilation_rate=(1, 1),
                 activation=None,
                 use_bias=True,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 kernel_regularizer=None,
                 bias_regularizer=None,
                 activity_regularizer=None,
                 kernel_constraint=None,
                 bias_constraint=None,
                 **kwargs):
        super(RProjConv2D, self).__init__(
            offset_creator_class=offset_creator_class,
            weight_basis=weight_basis,
            rank=2,
            filters=filters,
            kernel_size=kernel_size,
            strides=strides,
            padding=padding,
            data_format=data_format,
            dilation_rate=dilation_rate,
            activation=activation,
            use_bias=use_bias,
            kernel_initializer=kernel_initializer,
            bias_initializer=bias_initializer,
            kernel_regularizer=kernel_regularizer,
            bias_regularizer=bias_regularizer,
            activity_regularizer=activity_regularizer,
            kernel_constraint=kernel_constraint,
            bias_constraint=bias_constraint,
            **kwargs)
        self.input_spec = InputSpec(ndim=4)

    def build(self, input_shape):
        assert self.data_format != 'channels_first','only b01c supported'
        channel_axis = -1
        if input_shape[channel_axis] is None:
            raise ValueError('The channel dimension of the inputs '
                             'should be defined. Found `None`.')

        input_dim = input_shape[-1]

        self.units = self.filters

        kernel_shape = self.kernel_size + (input_dim, self.filters)
        self.kernel = self.add_weight(shape=kernel_shape,
                                      initializer=self.kernel_initializer,
                                      name='kernel',
                                      regularizer=self.kernel_regularizer,
                                      constraint=self.kernel_constraint)

        if self.use_bias:
            self.bias = self.add_weight(shape=(self.filters,),
                                        initializer=self.bias_initializer,
                                        name='bias',
                                        regularizer=self.bias_regularizer,
                                        constraint=self.bias_constraint)
        else:
            self.bias = None
        # Set input spec.
        self.input_spec = InputSpec(ndim=self.rank + 2,
                                    axes={channel_axis: input_dim})
        self.built = True

    def call(self, inputs):
        assert self.rank == 2, 'only conv2d supported for now...'

        kt0, eproj = self.kernel
        k = tf.add(kt0, eproj())

        bt0, eprojb = self.bias
        b = tf.add(bt0, eprojb())

        if self.rank == 2:
            outputs = K.conv2d(
                inputs,
                k,
                strides=self.strides,
                padding=self.padding,
                data_format=self.data_format,
                dilation_rate=self.dilation_rate)

        if self.use_bias:
            outputs = K.bias_add(
                outputs,
                b,
                data_format=self.data_format)


        #if self.activation is not None:
        #    assert False,'activation functions not supported'
        #    return self.activation(outputs)
        return outputs

    def compute_output_shape(self, input_shape):
        if self.data_format == 'channels_last':
            space = input_shape[1:-1]
            new_space = []
            for i in range(len(space)):
                new_dim = conv_utils.conv_output_length(
                    space[i],
                    self.kernel_size[i],
                    padding=self.padding,
                    stride=self.strides[i],
                    dilation=self.dilation_rate[i])
                new_space.append(new_dim)
            #self.filters*2 to accomodate LU representation
            return (input_shape[0],) + tuple(new_space) + (self.filters,)


class RProjBatchNormalization(LowRankBasisLayer):
    '''RProj version of BatchNormalization.'''

    def __init__(self, offset_creator_class, weight_basis,
                 axis=-1,
                 momentum=0.99,
                 epsilon=1e-3,
                 center=True,
                 scale=True,
                 beta_initializer='zeros',
                 gamma_initializer='ones',
                 moving_mean_initializer='zeros',
                 moving_variance_initializer='ones',
                 beta_regularizer=None,
                 gamma_regularizer=None,
                 beta_constraint=None,
                 gamma_constraint=None,
                 **kwargs):
        super(RProjBatchNormalization, self).__init__(offset_creator_class, weight_basis, **kwargs)
        self.supports_masking = True
        self.axis = axis
        self.momentum = momentum
        self.epsilon = epsilon
        self.center = center
        self.scale = scale
        self.beta_initializer = initializers.get(beta_initializer)
        self.gamma_initializer = initializers.get(gamma_initializer)
        self.moving_mean_initializer = initializers.get(moving_mean_initializer)
        self.moving_variance_initializer = initializers.get(moving_variance_initializer)
        self.beta_regularizer = regularizers.get(beta_regularizer)
        self.gamma_regularizer = regularizers.get(gamma_regularizer)
        self.beta_constraint = constraints.get(beta_constraint)
        self.gamma_constraint = constraints.get(gamma_constraint)

    def build(self, input_shape):
        dim = input_shape[self.axis]
        if dim is None:
            raise ValueError('Axis ' + str(self.axis) + ' of '
                             'input tensor should have a defined dimension '
                             'but the layer received an input with shape ' +
                             str(input_shape) + '.')
        self.input_spec = InputSpec(ndim=len(input_shape),
                                    axes={self.axis: dim})
        shape = (dim,)

        if self.scale:
            self.gamma = self.add_weight(shape=shape,
                                         name='gamma',
                                         initializer=self.gamma_initializer,
                                         regularizer=self.gamma_regularizer,
                                         constraint=self.gamma_constraint)
        else:
            self.gamma = None
        if self.center:
            self.beta = self.add_weight(shape=shape,
                                        name='beta',
                                        initializer=self.beta_initializer,
                                        regularizer=self.beta_regularizer,
                                        constraint=self.beta_constraint)
        else:
            self.beta = None
        self.moving_mean = self.add_non_trainable_weight(
            shape=shape,
            name='moving_mean',
            initializer=self.moving_mean_initializer)
        self.moving_variance = self.add_non_trainable_weight(
            shape=shape,
            name='moving_variance',
            initializer=self.moving_variance_initializer)
        self.built = True

    def call(self, inputs, training=None):
        input_shape = K.int_shape(inputs)
        # Prepare broadcasting shape.
        ndim = len(input_shape)
        reduction_axes = list(range(len(input_shape)))
        del reduction_axes[self.axis]
        broadcast_shape = [1] * len(input_shape)
        broadcast_shape[self.axis] = input_shape[self.axis]

        # Determines whether broadcasting is needed.
        needs_broadcasting = (sorted(reduction_axes) != list(range(ndim))[:-1])

        def normalize_inference():
            if needs_broadcasting:
                # In this case we must explicitly broadcast all parameters.
                broadcast_moving_mean = K.reshape(self.moving_mean,
                                                  broadcast_shape)
                broadcast_moving_variance = K.reshape(self.moving_variance,
                                                      broadcast_shape)
                if self.center:
                    broadcast_beta = K.reshape(self.beta, broadcast_shape)
                else:
                    broadcast_beta = None
                if self.scale:
                    broadcast_gamma = K.reshape(self.gamma,
                                                broadcast_shape)
                else:
                    broadcast_gamma = None
                return K.batch_normalization(
                    inputs,
                    broadcast_moving_mean,
                    broadcast_moving_variance,
                    broadcast_beta,
                    broadcast_gamma,
                    epsilon=self.epsilon)
            else:
                return K.batch_normalization(
                    inputs,
                    self.moving_mean,
                    self.moving_variance,
                    self.beta,
                    self.gamma,
                    epsilon=self.epsilon)

        # If the learning phase is *static* and set to inference:
        if training in {0, False}:
            return normalize_inference()

        # If the learning is either dynamic, or set to training:
        normed_training, mean, variance = K.normalize_batch_in_training(
            inputs, self.gamma, self.beta, reduction_axes,
            epsilon=self.epsilon)

        self.add_update([K.moving_average_update(self.moving_mean,
                                                 mean,
                                                 self.momentum),
                         K.moving_average_update(self.moving_variance,
                                                 variance,
                                                 self.momentum)],
                        inputs)

        # Pick the normalized form corresponding to the training phase.
        return K.in_train_phase(normed_training,
                                normalize_inference,
                                training=training)



class RProjLocallyConnected2D(LowRankBasisLayer):
    """Locally-connected layer for 2D inputs.
    The `LocallyConnected2D` layer works similarly
    to the `Conv2D` layer, except that weights are unshared,
    that is, a different set of filters is applied at each
    different patch of the input.
    """

    # @interfaces.legacy_conv2d_support
    def __init__(self, offset_creator_class, weight_basis,
                 filters,
                 kernel_size,
                 strides=(1, 1),
                 padding='valid',
                 data_format=None,
                 activation=None,
                 use_bias=True,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 kernel_regularizer=None,
                 bias_regularizer=None,
                 activity_regularizer=None,
                 kernel_constraint=None,
                 bias_constraint=None,
                 **kwargs):
        super(RProjLocallyConnected2D, self).__init__(offset_creator_class, weight_basis, **kwargs)
        self.filters = filters
        self.kernel_size = conv_utils.normalize_tuple(kernel_size, 2, 'kernel_size')
        self.strides = conv_utils.normalize_tuple(strides, 2, 'strides')
        self.padding = conv_utils.normalize_padding(padding)
        if self.padding != 'valid':
            raise ValueError('Invalid border mode for LocallyConnected2D '
                             '(only "valid" is supported): ' + padding)
        self.data_format = conv_utils.normalize_data_format(data_format)
        self.activation = activations.get(activation)
        self.use_bias = use_bias
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.bias_initializer = initializers.get(bias_initializer)
        self.kernel_regularizer = regularizers.get(kernel_regularizer)
        self.bias_regularizer = regularizers.get(bias_regularizer)
        self.activity_regularizer = regularizers.get(activity_regularizer)
        self.kernel_constraint = constraints.get(kernel_constraint)
        self.bias_constraint = constraints.get(bias_constraint)
        self.input_spec = InputSpec(ndim=4)

    def build(self, input_shape):
        if self.data_format == 'channels_last':
            input_row, input_col = input_shape[1:-1]
            input_filter = input_shape[3]
        else:
            input_row, input_col = input_shape[2:]
            input_filter = input_shape[1]
        if input_row is None or input_col is None:
            raise ValueError('The spatial dimensions of the inputs to '
                             ' a LocallyConnected2D layer '
                             'should be fully-defined, but layer received '
                             'the inputs shape ' + str(input_shape))
        output_row = conv_utils.conv_output_length(input_row, self.kernel_size[0],
                                                   self.padding, self.strides[0])
        output_col = conv_utils.conv_output_length(input_col, self.kernel_size[1],
                                                   self.padding, self.strides[1])
        self.output_row = output_row
        self.output_col = output_col
        self.kernel_shape = (output_row * output_col,
                             self.kernel_size[0] * self.kernel_size[1] * input_filter,
                             self.filters)
        self.kernel = self.add_weight(shape=self.kernel_shape,
                                      initializer=self.kernel_initializer,
                                      name='kernel',
                                      regularizer=self.kernel_regularizer,
                                      constraint=self.kernel_constraint)
        if self.use_bias:
            self.bias = self.add_weight(shape=(output_row, output_col, self.filters),
                                        initializer=self.bias_initializer,
                                        name='bias',
                                        regularizer=self.bias_regularizer,
                                        constraint=self.bias_constraint)
        else:
            self.bias = None
        if self.data_format == 'channels_first':
            self.input_spec = InputSpec(ndim=4, axes={1: input_filter})
        else:
            self.input_spec = InputSpec(ndim=4, axes={-1: input_filter})
        self.built = True

    def compute_output_shape(self, input_shape):
        if self.data_format == 'channels_first':
            rows = input_shape[2]
            cols = input_shape[3]
        elif self.data_format == 'channels_last':
            rows = input_shape[1]
            cols = input_shape[2]

        rows = conv_utils.conv_output_length(rows, self.kernel_size[0],
                                             self.padding, self.strides[0])
        cols = conv_utils.conv_output_length(cols, self.kernel_size[1],
                                             self.padding, self.strides[1])

        if self.data_format == 'channels_first':
            return (input_shape[0], self.filters, rows, cols)
        elif self.data_format == 'channels_last':
            return (input_shape[0], rows, cols, self.filters)

    def call(self, inputs):
        _, _, filters = self.kernel_shape

        output = K.local_conv2d(inputs,
                                self.kernel,
                                self.kernel_size,
                                self.strides,
                                (self.output_row, self.output_col),
                                self.data_format)

        if self.use_bias:
            if self.data_format == 'channels_first' or self.data_format == 'channels_last':
                output = K.bias_add(output, self.bias, data_format=self.data_format)

        output = self.activation(output)
        return output

    def get_config(self):
        config = {
            'filters': self.filters,
            'kernel_size': self.kernel_size,
            'strides': self.strides,
            'padding': self.padding,
            'data_format': self.data_format,
            'activation': activations.serialize(self.activation),
            'use_bias': self.use_bias,
            'kernel_initializer': initializers.serialize(self.kernel_initializer),
            'bias_initializer': initializers.serialize(self.bias_initializer),
            'kernel_regularizer': regularizers.serialize(self.kernel_regularizer),
            'bias_regularizer': regularizers.serialize(self.bias_regularizer),
            'activity_regularizer': regularizers.serialize(self.activity_regularizer),
            'kernel_constraint': constraints.serialize(self.kernel_constraint),
            'bias_constraint': constraints.serialize(self.bias_constraint)
        }
        base_config = super(LocallyConnected2D, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

tensorflow.keras.backend.floatx tensorflow.keras.constraints.serialize tensorflow.keras.regularizers.serialize tensorflow.keras.backend.moving_average_update tensorflow.keras.backend.batch_normalization tensorflow.python.keras.utils.conv_utils.conv_output_length tensorflow.keras.backend.normalize_batch_in_training tensorflow.keras.backend.local_conv2d tensorflow.keras.backend.in_train_phase tensorflow.keras.backend.int_shape tensorflow.keras.backend.reshape tensorflow.keras.layers.InputSpec tensorflow.keras.initializers.get tensorflow.keras.backend.bias_add tensorflow.keras.backend.conv3d tensorflow.keras.backend.dot tensorflow.keras.initializers.serialize tensorflow.python.keras.utils.conv_utils.normalize_tuple tensorflow.keras.backend.conv2d tensorflow.keras.constraints.get tensorflow.keras.activations.serialize tensorflow.keras.regularizers.get tensorflow.keras.backend.conv1d tensorflow.python.keras.utils.conv_utils.normalize_data_format tensorflow.keras.activations.get tensorflow.python.keras.utils.conv_utils.normalize_padding
[(78, 'tensorflow.keras.initializers.get', 'initializers.get', (['initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (130, 'tensorflow.keras.initializers.get', 'initializers.get', (['initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (161, 'tensorflow.keras.activations.get', 'activations.get', (['activation'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (163, 'tensorflow.keras.initializers.get', 'initializers.get', (['kernel_initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (164, 'tensorflow.keras.initializers.get', 'initializers.get', (['bias_initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (165, 'tensorflow.keras.regularizers.get', 'regularizers.get', (['kernel_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (166, 'tensorflow.keras.regularizers.get', 'regularizers.get', (['bias_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (167, 'tensorflow.keras.regularizers.get', 'regularizers.get', (['activity_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (168, 'tensorflow.keras.constraints.get', 'constraints.get', (['kernel_constraint'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (169, 'tensorflow.keras.constraints.get', 'constraints.get', (['bias_constraint'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (170, 'tensorflow.keras.layers.InputSpec', 'InputSpec', ([], {'min_ndim': '(2)'}), False, 'from tensorflow.keras.layers import Layer, InputSpec\n'), (190, 'tensorflow.keras.layers.InputSpec', 'InputSpec', ([], {'min_ndim': '(2)', 'axes': '{(-1): input_dim}'}), False, 'from tensorflow.keras.layers import Layer, InputSpec\n'), (202, 'tensorflow.keras.backend.dot', 'K.dot', (['inputs', 'k'], {}), True, 'import tensorflow.keras.backend as K\n'), (244, 'tensorflow.python.keras.utils.conv_utils.normalize_tuple', 'conv_utils.normalize_tuple', (['kernel_size', 'rank', '"""kernel_size"""'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (245, 'tensorflow.python.keras.utils.conv_utils.normalize_tuple', 'conv_utils.normalize_tuple', (['strides', 'rank', '"""strides"""'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (246, 'tensorflow.python.keras.utils.conv_utils.normalize_padding', 'conv_utils.normalize_padding', (['padding'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (247, 'tensorflow.python.keras.utils.conv_utils.normalize_data_format', 'conv_utils.normalize_data_format', (['data_format'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (248, 'tensorflow.python.keras.utils.conv_utils.normalize_tuple', 'conv_utils.normalize_tuple', (['dilation_rate', 'rank', '"""dilation_rate"""'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (249, 'tensorflow.keras.activations.get', 'activations.get', (['activation'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (251, 'tensorflow.keras.initializers.get', 'initializers.get', (['kernel_initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (252, 'tensorflow.keras.initializers.get', 'initializers.get', (['bias_initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (253, 'tensorflow.keras.regularizers.get', 'regularizers.get', (['kernel_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (254, 'tensorflow.keras.regularizers.get', 'regularizers.get', (['bias_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (255, 'tensorflow.keras.regularizers.get', 'regularizers.get', (['activity_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (256, 'tensorflow.keras.constraints.get', 'constraints.get', (['kernel_constraint'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (257, 'tensorflow.keras.constraints.get', 'constraints.get', (['bias_constraint'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (258, 'tensorflow.keras.layers.InputSpec', 'InputSpec', ([], {'ndim': '(self.rank + 2)'}), False, 'from tensorflow.keras.layers import Layer, InputSpec\n'), (285, 'tensorflow.keras.layers.InputSpec', 'InputSpec', ([], {'ndim': '(self.rank + 2)', 'axes': '{channel_axis: input_dim}'}), False, 'from tensorflow.keras.layers import Layer, InputSpec\n'), (395, 'tensorflow.keras.layers.InputSpec', 'InputSpec', ([], {'ndim': '(4)'}), False, 'from tensorflow.keras.layers import Layer, InputSpec\n'), (424, 'tensorflow.keras.layers.InputSpec', 'InputSpec', ([], {'ndim': '(self.rank + 2)', 'axes': '{channel_axis: input_dim}'}), False, 'from tensorflow.keras.layers import Layer, InputSpec\n'), (499, 'tensorflow.keras.initializers.get', 'initializers.get', (['beta_initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (500, 'tensorflow.keras.initializers.get', 'initializers.get', (['gamma_initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (501, 'tensorflow.keras.initializers.get', 'initializers.get', (['moving_mean_initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (502, 'tensorflow.keras.initializers.get', 'initializers.get', (['moving_variance_initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (503, 'tensorflow.keras.regularizers.get', 'regularizers.get', (['beta_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (504, 'tensorflow.keras.regularizers.get', 'regularizers.get', (['gamma_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (505, 'tensorflow.keras.constraints.get', 'constraints.get', (['beta_constraint'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (506, 'tensorflow.keras.constraints.get', 'constraints.get', (['gamma_constraint'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (546, 'tensorflow.keras.backend.int_shape', 'K.int_shape', (['inputs'], {}), True, 'import tensorflow.keras.backend as K\n'), (594, 'tensorflow.keras.backend.normalize_batch_in_training', 'K.normalize_batch_in_training', (['inputs', 'self.gamma', 'self.beta', 'reduction_axes'], {'epsilon': 'self.epsilon'}), True, 'import tensorflow.keras.backend as K\n'), (607, 'tensorflow.keras.backend.in_train_phase', 'K.in_train_phase', (['normed_training', 'normalize_inference'], {'training': 'training'}), True, 'import tensorflow.keras.backend as K\n'), (640, 'tensorflow.python.keras.utils.conv_utils.normalize_tuple', 'conv_utils.normalize_tuple', (['kernel_size', '(2)', '"""kernel_size"""'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (641, 'tensorflow.python.keras.utils.conv_utils.normalize_tuple', 'conv_utils.normalize_tuple', (['strides', '(2)', '"""strides"""'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (642, 'tensorflow.python.keras.utils.conv_utils.normalize_padding', 'conv_utils.normalize_padding', (['padding'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (646, 'tensorflow.python.keras.utils.conv_utils.normalize_data_format', 'conv_utils.normalize_data_format', (['data_format'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (647, 'tensorflow.keras.activations.get', 'activations.get', (['activation'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (649, 'tensorflow.keras.initializers.get', 'initializers.get', (['kernel_initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (650, 'tensorflow.keras.initializers.get', 'initializers.get', (['bias_initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (651, 'tensorflow.keras.regularizers.get', 'regularizers.get', (['kernel_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (652, 'tensorflow.keras.regularizers.get', 'regularizers.get', (['bias_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (653, 'tensorflow.keras.regularizers.get', 'regularizers.get', (['activity_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (654, 'tensorflow.keras.constraints.get', 'constraints.get', (['kernel_constraint'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (655, 'tensorflow.keras.constraints.get', 'constraints.get', (['bias_constraint'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (656, 'tensorflow.keras.layers.InputSpec', 'InputSpec', ([], {'ndim': '(4)'}), False, 'from tensorflow.keras.layers import Layer, InputSpec\n'), (670, 'tensorflow.python.keras.utils.conv_utils.conv_output_length', 'conv_utils.conv_output_length', (['input_row', 'self.kernel_size[0]', 'self.padding', 'self.strides[0]'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (672, 'tensorflow.python.keras.utils.conv_utils.conv_output_length', 'conv_utils.conv_output_length', (['input_col', 'self.kernel_size[1]', 'self.padding', 'self.strides[1]'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (706, 'tensorflow.python.keras.utils.conv_utils.conv_output_length', 'conv_utils.conv_output_length', (['rows', 'self.kernel_size[0]', 'self.padding', 'self.strides[0]'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (708, 'tensorflow.python.keras.utils.conv_utils.conv_output_length', 'conv_utils.conv_output_length', (['cols', 'self.kernel_size[1]', 'self.padding', 'self.strides[1]'], {}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (719, 'tensorflow.keras.backend.local_conv2d', 'K.local_conv2d', (['inputs', 'self.kernel', 'self.kernel_size', 'self.strides', '(self.output_row, self.output_col)', 'self.data_format'], {}), True, 'import tensorflow.keras.backend as K\n'), (80, 'tensorflow.keras.backend.floatx', 'K.floatx', ([], {}), True, 'import tensorflow.keras.backend as K\n'), (132, 'tensorflow.keras.backend.floatx', 'K.floatx', ([], {}), True, 'import tensorflow.keras.backend as K\n'), (204, 'tensorflow.keras.backend.bias_add', 'K.bias_add', (['output', 'b'], {}), True, 'import tensorflow.keras.backend as K\n'), (291, 'tensorflow.keras.backend.conv1d', 'K.conv1d', (['inputs', 'self.kernel'], {'strides': 'self.strides[0]', 'padding': 'self.padding', 'data_format': 'self.data_format', 'dilation_rate': 'self.dilation_rate[0]'}), True, 'import tensorflow.keras.backend as K\n'), (299, 'tensorflow.keras.backend.conv2d', 'K.conv2d', (['inputs', 'self.kernel'], {'strides': 'self.strides', 'padding': 'self.padding', 'data_format': 'self.data_format', 'dilation_rate': 'self.dilation_rate'}), True, 'import tensorflow.keras.backend as K\n'), (307, 'tensorflow.keras.backend.conv3d', 'K.conv3d', (['inputs', 'self.kernel'], {'strides': 'self.strides', 'padding': 'self.padding', 'data_format': 'self.data_format', 'dilation_rate': 'self.dilation_rate'}), True, 'import tensorflow.keras.backend as K\n'), (316, 'tensorflow.keras.backend.bias_add', 'K.bias_add', (['outputs', 'self.bias'], {'data_format': 'self.data_format'}), True, 'import tensorflow.keras.backend as K\n'), (438, 'tensorflow.keras.backend.conv2d', 'K.conv2d', (['inputs', 'k'], {'strides': 'self.strides', 'padding': 'self.padding', 'data_format': 'self.data_format', 'dilation_rate': 'self.dilation_rate'}), True, 'import tensorflow.keras.backend as K\n'), (447, 'tensorflow.keras.backend.bias_add', 'K.bias_add', (['outputs', 'b'], {'data_format': 'self.data_format'}), True, 'import tensorflow.keras.backend as K\n'), (693, 'tensorflow.keras.layers.InputSpec', 'InputSpec', ([], {'ndim': '(4)', 'axes': '{(1): input_filter}'}), False, 'from tensorflow.keras.layers import Layer, InputSpec\n'), (695, 'tensorflow.keras.layers.InputSpec', 'InputSpec', ([], {'ndim': '(4)', 'axes': '{(-1): input_filter}'}), False, 'from tensorflow.keras.layers import Layer, InputSpec\n'), (740, 'tensorflow.keras.activations.serialize', 'activations.serialize', (['self.activation'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (742, 'tensorflow.keras.initializers.serialize', 'initializers.serialize', (['self.kernel_initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (743, 'tensorflow.keras.initializers.serialize', 'initializers.serialize', (['self.bias_initializer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (744, 'tensorflow.keras.regularizers.serialize', 'regularizers.serialize', (['self.kernel_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (745, 'tensorflow.keras.regularizers.serialize', 'regularizers.serialize', (['self.bias_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (746, 'tensorflow.keras.regularizers.serialize', 'regularizers.serialize', (['self.activity_regularizer'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (747, 'tensorflow.keras.constraints.serialize', 'constraints.serialize', (['self.kernel_constraint'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (748, 'tensorflow.keras.constraints.serialize', 'constraints.serialize', (['self.bias_constraint'], {}), False, 'from tensorflow.keras import regularizers, constraints, initializers, activations\n'), (84, 'src.subspace.keras_ext.rproj_layers_util._convert_string_dtype', '_convert_string_dtype', (['dtype'], {}), False, 'from src.subspace.keras_ext.rproj_layers_util import _convert_string_dtype\n'), (330, 'tensorflow.python.keras.utils.conv_utils.conv_output_length', 'conv_utils.conv_output_length', (['space[i]', 'self.kernel_size[i]'], {'padding': 'self.padding', 'stride': 'self.strides[i]', 'dilation': 'self.dilation_rate[i]'}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (342, 'tensorflow.python.keras.utils.conv_utils.conv_output_length', 'conv_utils.conv_output_length', (['space[i]', 'self.kernel_size[i]'], {'padding': 'self.padding', 'stride': 'self.strides[i]', 'dilation': 'self.dilation_rate[i]'}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (463, 'tensorflow.python.keras.utils.conv_utils.conv_output_length', 'conv_utils.conv_output_length', (['space[i]', 'self.kernel_size[i]'], {'padding': 'self.padding', 'stride': 'self.strides[i]', 'dilation': 'self.dilation_rate[i]'}), False, 'from tensorflow.python.keras.utils import conv_utils\n'), (560, 'tensorflow.keras.backend.reshape', 'K.reshape', (['self.moving_mean', 'broadcast_shape'], {}), True, 'import tensorflow.keras.backend as K\n'), (562, 'tensorflow.keras.backend.reshape', 'K.reshape', (['self.moving_variance', 'broadcast_shape'], {}), True, 'import tensorflow.keras.backend as K\n'), (573, 'tensorflow.keras.backend.batch_normalization', 'K.batch_normalization', (['inputs', 'broadcast_moving_mean', 'broadcast_moving_variance', 'broadcast_beta', 'broadcast_gamma'], {'epsilon': 'self.epsilon'}), True, 'import tensorflow.keras.backend as K\n'), (581, 'tensorflow.keras.backend.batch_normalization', 'K.batch_normalization', (['inputs', 'self.moving_mean', 'self.moving_variance', 'self.beta', 'self.gamma'], {'epsilon': 'self.epsilon'}), True, 'import tensorflow.keras.backend as K\n'), (598, 'tensorflow.keras.backend.moving_average_update', 'K.moving_average_update', (['self.moving_mean', 'mean', 'self.momentum'], {}), True, 'import tensorflow.keras.backend as K\n'), (601, 'tensorflow.keras.backend.moving_average_update', 'K.moving_average_update', (['self.moving_variance', 'variance', 'self.momentum'], {}), True, 'import tensorflow.keras.backend as K\n'), (728, 'tensorflow.keras.backend.bias_add', 'K.bias_add', (['output', 'self.bias'], {'data_format': 'self.data_format'}), True, 'import tensorflow.keras.backend as K\n'), (565, 'tensorflow.keras.backend.reshape', 'K.reshape', (['self.beta', 'broadcast_shape'], {}), True, 'import tensorflow.keras.backend as K\n'), (569, 'tensorflow.keras.backend.reshape', 'K.reshape', (['self.gamma', 'broadcast_shape'], {}), True, 'import tensorflow.keras.backend as K\n')]