"""Contains custom downsampling classes."""

import tensorflow as tf

from layers.group_norm import GroupNormalization

def get_downsampling(downsampling):

    if downsampling == 'max':

        return MaxDownsample

    elif downsampling == 'conv':

        return ConvDownsample

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

    def __init__(self,

                 filters,

                 data_format='channels_last',

                 groups=8,

                 l2_scale=1e-5,

                 **kwargs):

        super(ConvDownsample, self).__init__()

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

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

                            'data_format': data_format,

                            'groups': groups,

                            'l2_scale': l2_scale})

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

                            filters=filters,

                            kernel_size=3,

                            strides=2,

                            padding='same',

                            data_format=data_format,

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

                            kernel_initializer='he_normal')

        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 MaxDownsample(tf.keras.layers.Layer):

    def __init__(self,

                 data_format='channels_last',

                 **kwargs):

        super(MaxDownsample, self).__init__()

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

        self.config.update({'data_format': data_format})

        self.maxpool = tf.keras.layers.MaxPooling3D(

                            pool_size=2,

                            strides=2,

                            padding='same',

                            data_format=data_format)

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

        inputs = self.maxpool(inputs)

        return inputs

    def get_config(self):

        return self.config