import tensorflow as tf

from tensorflow.keras import models, backend

from tensorflow.keras.layers import Input

from . import functional


class ShallowACRBM(models.Model):

    def __init__(self, input_shape, hidden_dim, visible, *args, **kwargs):

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

        self.hidden_dim = hidden_dim
        self.visible = visible

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

        self.reconstruction_loss = kwargs.pop('reconstruction_loss', functional.zero_loss)

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

        self.visible_states = Input(shape=input_shape)

        self.acrbm_outputs = functional.acrbm(self.visible_states, units=self.hidden_dim,
                                              visible=self.visible,  *args, **kwargs)

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

        super(ShallowACRBM, self).__init__(inputs=self.visible_states, outputs=self.acrbm_outputs)

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

    def compile(self, optimizer='sgd', **kwargs):

        super(ShallowACRBM, self).compile(optimizer=optimizer, **kwargs)

    def get_config(self):

        return super(ShallowACRBM, self).get_config()

    def call(self, x, training=None, mask=None):

        return super(ShallowACRBM, self).call(x, training=training, mask=mask)

    def train_step(self, data):

        if len(data) == 2:

            x, y = data

        else:

            x, y = data, None

        # forward pass
        with tf.GradientTape() as tape:

            outputs = self(x, training=True)

            loss = functional.partition_loss(outputs[0])

            if self.visible:

                loss += self.reconstruction_loss(x, outputs[1])

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

        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

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

        # update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

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

        if y is not None:

            y_pred = backend.argmax(outputs[0], axis=-1)

            self.compiled_metrics.update_state(y, y_pred)

        ret = {'loss': loss}

        ret.update({m.name: m.result() for m in self.metrics})

        return ret

    def test_step(self, data):

        if len(data) == 2:

            x, y = data

        else:

            x, y = data, None

        outputs = self(x, training=False)

        loss = functional.partition_loss(outputs[0])

        if self.visible:

            loss += self.reconstruction_loss(x, outputs[1])

        if y is not None:

            y_pred = backend.argmax(outputs[0], axis=-1)

            self.compiled_metrics.update_state(y, y_pred)

        ret = {'loss': loss}

        ret.update({m.name: m.result() for m in self.metrics})

        return ret

    def predict(self, x, **kwargs):

        outputs = super(ShallowACRBM, self).predict(x, **kwargs)

        if self.visible:

            return outputs

        else:

            return outputs[0]

    def predict_visible(self, h):

        return self.get_layer(self.layers[1].name).compute_visible(h)


class DeepACRBM(models.Sequential):

    def __init__(self, *args, **kwargs):

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

        super(DeepACRBM, self).__init__(*args, **kwargs)

    def stack(self, input_shape, hidden_dim, compile_callback=None, **kwargs):

        if kwargs.pop('visible', False):

            raise ValueError('Intermediate visible layers are not supported')

        layer = ShallowACRBM(input_shape=input_shape, hidden_dim=hidden_dim, visible=False, **kwargs)

        if compile_callback is not None:

            compile_callback(layer)

        return super(DeepACRBM, self).add(layer)

    def compile(self, optimizer='sgd', **kwargs):

        super(DeepACRBM, self).compile(optimizer=optimizer, **kwargs)

    def get_config(self):

        return super(DeepACRBM, self).get_config()

    def call(self, x, training=None, mask=None):

        outputs = super(DeepACRBM, self).call(x, training=training, mask=mask)

        if not isinstance(outputs, tuple):

            return outputs,

        return outputs

    def train_step(self, data):

        if len(data) == 2:

            h, y = data

        else:

            h, y = data, None

        # forward pass
        with tf.GradientTape() as tape:

            loss = 0.0

            for i in range(len(self.layers)):

                if isinstance(self.layers[i], ShallowACRBM):

                    state = (h, )

                    if y is not None:

                        state += (y, )

                    ret = self.layers[i].train_step(state)

                    loss += ret['loss']

                    h = self.layers[i](h)[0]

                else:

                    h = self.layers[i](h)

        if len(self.layers) > 1:

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

            trainable_vars = self.trainable_variables
            gradients = tape.gradient(loss, trainable_vars)

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

            # update weights
            self.optimizer.apply_gradients(zip(gradients, trainable_vars))

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

        ret['loss'] = loss

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

        return ret

    def predict(self, x, **kwargs):

        outputs = super(DeepACRBM, self).predict(x, **kwargs)

        if len(self.layers) > 1:

            return outputs[0]

        else:

            return outputs

    def predict_visible_step(self, x, start=None, end=None):

        if start is None:

            start = 0

        if end is None:

            end = len(self.layers)

        if start >= end or end > len(self.layers):

            raise ValueError('...')

        h = x

        for i in range(start, end):

            if not hasattr(self.layers[i], 'predict_visible'):

                raise ValueError('...')

            h = self.layers[i](h, training=False)[0]

        for i in reversed(range(start, end)):

            h = self.layers[i].predict_visible(h)

        return h

    def predict_visible(self, x, batch_size=32, start=None, end=None):

        data_iter = tf.data.Dataset.from_tensor_slices(x).batch(batch_size).as_numpy_iterator()

        xi = data_iter.next()

        outputs = self.predict_visible_step(xi, start, end)

        for xi in data_iter:

            vi = self.predict_visible_step(xi, start, end)

            outputs = tf.concat([outputs, vi], axis=0)

        return outputs


class AEBasedACRBM(models.Model):

    def __init__(self, input_shape, encoder, acrbm, merge, decoder, feedback=True, *args, **kwargs):

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

        self.feedback = feedback

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

        inputs = Input(shape=input_shape)

        latent = encoder(inputs)

        h, = acrbm(latent)

        z = merge([latent, h])

        outputs = decoder(z)

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

        super(AEBasedACRBM, self).__init__(inputs=inputs, outputs=[outputs, latent, h], *args, **kwargs)

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

    def compile(self, optimizer='adam', **kwargs):

        super(AEBasedACRBM, self).compile(optimizer=optimizer, **kwargs)

    def get_config(self):

        return super(AEBasedACRBM, self).get_config()

    def call(self, x, training=None, mask=None):

        return super(AEBasedACRBM, self).call(x, training=training, mask=mask)

    def train_step(self, data):

        if len(data) == 2:

            x, y = data

        else:

            x, y = data, None

        # forward pass
        with tf.GradientTape() as tape:

            loss = 0.0

            inputs = self.layers[0](x)

            latent = self.layers[1](inputs, training=True)
            h, = self.layers[2](latent, training=True)

            merge = self.layers[3]([latent, h])

            outputs = self.layers[4](merge, training=True)

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

            if self.feedback:

                loss += self.layers[2].train_step((latent, ))['loss']

            else:

                loss += functional.partition_loss(h)

            loss += functional.squared_l2_loss(x, outputs)

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

        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

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

        # update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

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

        if y is not None:

            y_pred = backend.argmax(h, axis=-1)

            self.compiled_metrics.update_state(y, y_pred)

        ret = {'loss': loss}

        ret.update({m.name: m.result() for m in self.metrics})

        return ret

    def predict_visible_step(self, x):

        inputs = self.layers[0](x)

        latent = self.layers[1](inputs)
        h = self.layers[2](latent)

        v = self.layers[2].predict_visible_step(h)

        return v

    def predict_visible(self, x, batch_size=32):

        data_iter = tf.data.Dataset.from_tensor_slices(x).batch(batch_size).as_numpy_iterator()

        xi = data_iter.next()

        v = self.predict_visible_step(xi)

        for xi in data_iter:

            vi = self.predict_visible_step(xi)

            v = tf.concat([v, vi], axis=0)

        return v
