import tensorflow as tf

from tensorflow.keras import models

from tensorflow.keras.layers import Input

from . import functional


class AutoSegmentationBaseline(models.Model):

    def __init__(self, input_shape, encoder, loc2d, rho, encoder_loss=None, *args, **kwargs):

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

        inputs = Input(shape=input_shape)

        latent = encoder(inputs)

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

        outputs = loc2d([inputs, latent])

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

        super(AutoSegmentationBaseline, self).__init__(inputs=inputs, outputs=[outputs], *args, **kwargs)

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

        self.g_loss = functional.g_loss(rho)

        if encoder_loss is None:

            self.encoder_loss = functional.zero_loss

        else:

            self.encoder_loss = encoder_loss

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

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

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

    def get_config(self):

        return super(AutoSegmentationBaseline, self).get_config()

    def reassign_rho(self, rho):

        self.g_loss = functional.g_loss(rho)

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

        return super(AutoSegmentationBaseline, 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:

            inputs = self.layers[0](x)
            latent = self.layers[1](inputs, training=True)
            outputs = self.layers[2]([inputs, latent], training=True)

            loss = self.encoder_loss(latent)
            loss += self.g_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:

            self.compiled_metrics.update_state(y, outputs)

        ret = {'loss': loss}

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

        return ret
