from copy import deepcopy

import tensorflow as tf

from tensorflow.keras import activations
from tensorflow.keras.layers import Layer


class ACRBM(Layer):

    def __init__(self, units, visible: bool, *args, **kwargs):

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

        copy_kwargs = deepcopy(kwargs)

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

        self.units = units

        self.visible = visible

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

        self.sigma_h = activations.get(copy_kwargs.pop('sigma_h', 'softmax'))
        self.sigma_v = activations.get(copy_kwargs.pop('sigma_v', 'sigmoid'))

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

        self.use_bh = copy_kwargs.pop('use_bh', True)
        self.use_bv = copy_kwargs.pop('use_bv', False)

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

        self.bh_trainable = copy_kwargs.pop('bh_trainable', True)
        self.bv_trainable = copy_kwargs.pop('bv_trainable', True)

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

        self.kernel_init = copy_kwargs.pop('kernel_initializer', 'glorot_uniform')

        self.bh_init = copy_kwargs.pop('bh_initializer', 'zeros')
        self.bv_init = copy_kwargs.pop('bv_initializer', 'zeros')

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

        self.kernel_regularizer = copy_kwargs.pop('kernel_regularizer', None)

        self.bh_regularizer = copy_kwargs.pop('bh_regularizer', None)
        self.bv_regularizer = copy_kwargs.pop('bv_regularizer', None)

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

        self.kernel_constraint = copy_kwargs.pop('kernel_constraint', None)

        self.bh_constraint = copy_kwargs.pop('bh_constraint', None)
        self.bv_constraint = copy_kwargs.pop('bv_constraint', None)

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

        self.w = None
        self.bh = None
        self.bv = None

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

        super(ACRBM, self).__init__(*args, **copy_kwargs)

    def validate_attributes_consistency(self, inputs_shape):

        pass

    def init_kernel(self, inputs_shape):

        kernel_shape = (inputs_shape[-1], self.units)

        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_visible(self, inputs_shape):

        if self.use_bv:

            bv_shape = (1, inputs_shape[-1])

            self.bv = self.add_weight(name='bv', shape=bv_shape, dtype=tf.float32,
                                      initializer=self.bv_init, regularizer=self.bv_regularizer,
                                      trainable=self.bv_trainable, constraint=self.bv_constraint)

    def init_hidden(self):

        if self.use_bh:

            bh_shape = (1, self.units)

            self.bh = self.add_weight(name='bh', shape=bh_shape, dtype=tf.float32,
                                      initializer=self.bh_init, regularizer=self.bh_regularizer,
                                      trainable=self.bh_trainable, constraint=self.bh_constraint)

    def compute_hidden(self, inputs):

        h = tf.matmul(inputs, self.w)

        if self.use_bh:
            h += self.bh

        h = self.sigma_h(h)

        return h

    def compute_visible(self, h):

        v = tf.matmul(h, tf.transpose(self.w))

        if self.use_bv:

            v += self.bv

        v = self.sigma_v(v)

        return v

    def build(self, inputs_shape):

        self.validate_attributes_consistency(inputs_shape)

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

        self.init_kernel(inputs_shape)
        self.init_hidden()

        if self.visible:

            self.init_visible(inputs_shape)

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

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

        h = self.compute_hidden(inputs)

        if self.visible:

            v = self.compute_visible(h)

            return h, v

        else:

            return h,

    def get_config(self):

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

        return cfg
