"""Contains custom variational autoencoder class."""

import tensorflow as tf

from layers.downsample import get_downsampling

from layers.upsample import get_upsampling

from layers.resnet import ResnetBlock

def sample(inputs):

    """Samples from the Gaussian given by mean and variance."""

    z_mean, z_logvar = inputs

    eps = tf.random.normal(shape=z_mean.shape, dtype=tf.float32)

    return z_mean + tf.math.exp(0.5 * z_logvar) * eps

class VariationalAutoencoder(tf.keras.layers.Layer):

    def __init__(self,

                 data_format='channels_last',

                 groups=8,

                 reduction=2,

                 l2_scale=1e-5,

                 downsampling='conv',

                 upsampling='conv',

                 base_filters=16,

                 depth=4,

                 out_ch=2):

        """ Initializes the variational autoencoder: consists of sampling

            then an alternating series of SENet blocks and upsampling.

            References:

                - [3D MRI brain tumor segmentation using autoencoder regularization](https://arxiv.org/pdf/1810.11654.pdf)

        """

        super(VariationalAutoencoder, self).__init__()

        # Set up config for self.get_config() to serialize later.

        self.data_format = data_format

        self.l2_scale = l2_scale

        self.config = super(VariationalAutoencoder, self).get_config()

        self.config.update({'groups': groups,

                            'reduction': reduction,

                            'downsampling': downsampling,

                            'upsampling': upsampling,

                            'base_filters': base_filters,

                            'depth': depth,

                            'out_ch': out_ch})

        # Retrieve downsampling method.

        Downsample = get_downsampling(downsampling)

        # Retrieve upsampling method.

        Upsample = get_upsampling(upsampling)

        # Extra downsampling layer to reduce parameters.

        self.downsample = Downsample(

                            filters=base_filters//2,

                            groups=groups,

                            data_format=data_format,

                            kernel_regularizer=tf.keras.regularizers.l2(l=l2_scale))

        # Build sampling layers.

        self.flatten = tf.keras.layers.Flatten(data_format)

        self.proj = tf.keras.layers.Dense(

                            units=base_filters*(2**(depth-1)),

                            kernel_regularizer=tf.keras.regularizers.l2(l=l2_scale),

                            kernel_initializer='he_normal')

        self.latent_size = base_filters*(2**(depth-2))

        self.sample = tf.keras.layers.Lambda(sample)

        # Extra upsampling layer to counter extra downsampling layer.

        self.upsample = Upsample(

                            filters=base_filters*(2**(depth-1)),

                            groups=groups,

                            data_format=data_format,

                            l2_scale=l2_scale)

        # Build layers at all spatial levels.

        self.levels = []

        for i in range(depth-2, -1, -1):

            upsample = Upsample(

                        filters=base_filters*(2**i),

                        groups=groups,

                        data_format=data_format,

                        l2_scale=l2_scale)

            conv = ResnetBlock(

                        filters=base_filters*(2**i),

                        groups=groups,

                        reduction=reduction,

                        data_format=data_format,

                        l2_scale=l2_scale)

            self.levels.append([upsample, conv])

        # Output layer convolution.

        self.out = tf.keras.layers.Conv3D(

                            filters=out_ch,

                            kernel_size=3,

                            strides=1,

                            padding='same',

                            data_format=data_format,

                            kernel_regularizer=tf.keras.regularizers.l2(l=l2_scale),

                            kernel_initializer='he_normal')

    def build(self, input_shape):

        h, w, d = input_shape[1:-1] if self.data_format == 'channels_last' else input_shape[2:]

        # Build reshaping layers after sampling.

        self.unproj = tf.keras.layers.Dense(

                                units=h*w*d*1//8,

                                kernel_regularizer=tf.keras.regularizers.l2(l=self.l2_scale),

                                kernel_initializer='he_normal',

                                activation='relu')

        self.unflatten = tf.keras.layers.Reshape(

                                (h//2, w//2, d//2, 1) if self.data_format == 'channels_last' else (1, h//2, w//2, d//2))

    def call(self, inputs, training=None):

        # Downsample.

        inputs = self.downsample(inputs)

        # Flatten and project

        inputs = self.flatten(inputs)

        inputs = self.proj(inputs)

        # Sample.

        z_mean = inputs[:, :self.latent_size]

        z_logvar = inputs[:, self.latent_size:]

        inputs = self.sample([z_mean, z_logvar])

        # Restored projection and reshape

        inputs = self.unproj(inputs)

        inputs = self.unflatten(inputs)

        # Upsample.

        inputs = self.upsample(inputs)

        # Iterate through spatial levels.

        for level in self.levels:

            upsample, conv = level

            inputs = upsample(inputs, training=training)

            inputs = conv(inputs, training=training)

        # Map convolution to number of original input channels.

        inputs = self.out(inputs)

        return inputs, z_mean, z_logvar

    def get_config(self):

        self.config.update({'data_format': self.data_format,

                            'l2_scale': self.l2_scale})

        return self.config