import tensorflow as tf

from tensorflow.keras import backend

from tensorflow.keras.layers import Layer


class LOC2D(Layer):

    def __init__(self, strides=(1, 1), padding='same', *args, **kwargs):

        # -------------------------------------------

        self.kernel_init = kwargs.pop('kernel_initializer', 'ones')

        self.bias_init = kwargs.pop('bias_initializer', 'zeros')

        # -------------------------------------------

        self.kernel_regularizer = kwargs.pop('kernel_regularizer', None)
        self.bias_regularizer = kwargs.pop('bias_regularizer', None)

        # -------------------------------------------

        self.kernel_constraint = kwargs.pop('kernel_constraint', None)
        self.bias_constraint = kwargs.pop('bias_constraint', None)

        # -------------------------------------------

        self.use_bias = kwargs.pop('use_bias', False)

        # -------------------------------------------

        self.strides = strides
        self.padding = padding

        # -------------------------------------------

        self.w = None
        self.bias = None

        # -------------------------------------------

        super(LOC2D, self).__init__(trainable=kwargs.pop('trainable', False), *args, **kwargs)

        # -------------------------------------------

    def init_kernel(self, inputs_shape):

        in_shape = inputs_shape[0]
        z_shape = inputs_shape[1]

        kernel_shape = (1, 1, 1, in_shape[-1], z_shape[-1])

        self.w = self.add_weight(name='w', shape=kernel_shape, dtype=tf.float32,
                                 initializer=self.kernel_init, regularizer=self.kernel_regularizer,
                                 trainable=self.trainable, constraint=self.kernel_constraint)

    def init_bias(self, inputs_shape):

        if self.use_bias:

            z_shape = inputs_shape[1]

            bias_shape = (1, 1, 1, 1, z_shape[-1])

            self.bias = self.add_weight(name='b', shape=bias_shape, dtype=tf.float32,
                                        initializer=self.bias_init, regularizer=self.bias_regularizer,
                                        trainable=self.trainable, constraint=self.bias_constraint)

    def build(self, inputs_shape):

        self.init_kernel(inputs_shape)
        self.init_bias(inputs_shape)

    def nd_map(self, elems):

        xi = elems[0][None, ...]
        gi = elems[1]

        outi = backend.conv2d(x=xi, kernel=gi, strides=self.strides, padding=self.padding)

        return outi[0], 0.0

    def compute_observer(self, local_kernel):

        observer_kernel = local_kernel[..., None, :] * self.w

        if self.use_bias:

            observer_kernel += self.bias

        return observer_kernel

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

        observer_kernel = self.compute_observer(inputs[1])

        out, _ = tf.map_fn(self.nd_map, (inputs[0], observer_kernel))

        return out

    def get_config(self):

        cfg = super(LOC2D, self).get_config()

        return cfg
