"""Contains custom upsampling classes."""
import tensorflow as tf

from layers.group_norm import GroupNormalization


def get_upsampling(upsampling):
    if upsampling == 'linear':
        return LinearUpsample
    elif upsampling == 'conv':
        return ConvUpsample


class ConvUpsample(tf.keras.layers.Layer):
    def __init__(self,
                 filters,
                 groups=8,
                 data_format='channels_last',
                 l2_scale=1e-5,
                 **kwargs):
        super(ConvUpsample, self).__init__()
        self.config = super(ConvUpsample, self).get_config()
        self.config.update({'filters': filters,
                            'data_format': data_format,
                            'groups': groups,
                            'l2_scale': l2_scale})

        self.conv = tf.keras.layers.Conv3DTranspose(
                            filters=filters,
                            kernel_size=3,
                            strides=2,
                            padding='same',
                            data_format=data_format)
        self.norm = GroupNormalization(
                            groups=groups,
                            axis=-1 if data_format == 'channels_last' else 1)
        self.relu = tf.keras.layers.Activation('relu')

    def __call__(self, inputs, training=None):
        inputs = self.conv(inputs)
        inputs = self.norm(inputs, training=training)
        inputs = self.relu(inputs)
        return inputs

    def get_config(self):
        return self.config


class LinearUpsample(tf.keras.layers.Layer):
    def __init__(self,
                 filters,
                 data_format='channels_last',
                 l2_scale=1e-5,
                 **kwargs):
        super(LinearUpsample, self).__init__()
        self.config = super(LinearUpsample, self).get_config()
        self.config.update({'filters': filters,
                            'data_format': data_format,
                            'l2_scale': l2_scale})

        self.ptwise = tf.keras.layers.Conv3D(
                                filters=filters,
                                kernel_size=1,
                                strides=1,
                                padding='same',
                                data_format=data_format,
                                kernel_regularizer=tf.keras.regularizers.l2(l=l2_scale),
                                kernel_initializer='he_normal')
        self.linear = tf.keras.layers.UpSampling3D(
                                size=2,
                                data_format=data_format)

    def __call__(self, inputs, training=None):
        inputs = self.ptwise(inputs)
        inputs = self.linear(inputs)
        return inputs

    def get_config(self):
        return self.config