import tensorflow as tf

import tensorflow.keras.layers as tfkl

class Deeplabv3_plus(tf.keras.Model):

    def __init__(self,

                 num_classes,

                 kernel_size_initial_conv,

                 num_channels_atrous=512,

                 num_channels_DCNN=[256, 512, 1024],

                 num_channels_ASPP=256,

                 kernel_size_atrous=3,

                 kernel_size_DCNN=[1, 3],

                 kernel_size_ASPP=[1, 3, 3, 3],

                 num_filters_final_encoder=512,

                 num_channels_from_backcone=[128, 96],

                 num_channels_UpConv=[512, 256, 128],

                 kernel_size_UpConv=3,

                 stride_UpConv=(2, 2),

                 use_batchnorm_UpConv=False,

                 use_transpose_UpConv=False,

                 padding='same',

                 nonlinearity='relu',

                 use_batchnorm=True,

                 use_bias=True,

                 data_format='channels_last',

                 MultiGrid=[1, 2, 4],

                 rate_ASPP=[1, 6, 12, 18],

                 atrous_output_stride=16,

                 # Not adapted code for any other out stride

                 **kwargs):

        """ Arguments:

            kernel_size_initial_conv: the size of the kernel for the

                                      first convolution

            num_channels_DCNN: touple with the number of channels for the

                               first three blocks of the DCNN

            kernel_size_DCNN: two element touple with the kernel size of the

                              first and last convolution of the resnet_block

                              (First element) and the middle convolution

                              of the resnet_block (Second element)  """

        super(Deeplabv3_plus, self).__init__(**kwargs)

        self.num_classes = num_classes

        #ResNet backbone

        self.first_conv = tfkl.Conv2D(num_channels_DCNN[0],

                                      kernel_size_initial_conv,

                                      strides=2,

                                      padding=padding,

                                      use_bias=use_bias,

                                      data_format=data_format)

        self.block1 = resnet_block(False,

                                   num_channels_DCNN[0],

                                   kernel_size_DCNN,

                                   padding,

                                   nonlinearity,

                                   use_batchnorm,

                                   use_bias,

                                   data_format)

        self.block2 = resnet_block(True,

                                   num_channels_DCNN[1],

                                   kernel_size_DCNN,

                                   padding,

                                   nonlinearity,

                                   use_batchnorm,

                                   use_bias,

                                   data_format)

        self.block3 = resnet_block(True,

                                   num_channels_DCNN[2],

                                   kernel_size_DCNN,

                                   padding,

                                   nonlinearity,

                                   use_batchnorm,

                                   use_bias,

                                   data_format)

        #Atrous components

        self.atrous_conv = Atrous_conv(num_channels_atrous,

                                       kernel_size_atrous,

                                       MultiGrid,

                                       padding,

                                       use_batchnorm,

                                       'linear',

                                       use_bias,

                                       data_format,

                                       atrous_output_stride)

        self.aspp_term = atrous_spatial_pyramid_pooling(num_channels_ASPP,

                                                        kernel_size_ASPP,

                                                        rate_ASPP,

                                                        padding,

                                                        use_batchnorm,

                                                        'linear',

                                                        use_bias,

                                                        data_format)

        #Final convolution of encoder

        self.final_encoder_conv = aspp_block(1,

                                             1,

                                             num_filters_final_encoder,

                                             padding,

                                             use_batchnorm,

                                             'linear',

                                             use_bias,

                                             data_format)

        #Decoder

        self.decoder_term = Decoder(num_channels_from_backcone=num_channels_from_backcone,

                                    num_channels_UpConv=num_channels_UpConv,

                                    kernel_size_UpConv=kernel_size_UpConv,

                                    stride_UpConv=stride_UpConv,

                                    use_batchnorm_UpConv=use_batchnorm_UpConv,

                                    use_transpose_UpConv=use_transpose_UpConv,

                                    use_bias=use_bias,

                                    padding=padding,

                                    data_format=data_format)

        self.output_conv = tfkl.Conv2D(num_classes,

                                       1,

                                       activation='linear',

                                       padding='same',

                                       data_format=data_format)

    def call(self, x, training=False):

        ###Encoder

        before_final_stride = self.first_conv(x, training=training)  # output stride 2

        before_final_stride = self.block1(before_final_stride, training=training)  # output stride 2

        before_final_stride = self.block2(before_final_stride, training=training)  # output stride 4

        atrous_out = self.block3(before_final_stride, training=training)  # output stride 8

        atrous_out = self.atrous_conv(atrous_out, training=training)

        out = self.aspp_term(atrous_out, training=training)

        out = self.final_encoder_conv(out, training=training)

        ###Decoder

        out = self.decoder_term(atrous_out ,out, before_final_stride, training=training)

        out = self.output_conv(out, training=training)

        if self.num_classes == 1:

            out = tfkl.Activation('sigmoid')(out)

        else:

            out = tfkl.Activation('softmax')(out)

        # Upsample to same size as the input

        # print(f"Input Shape: {x.shape}, Out Shape: {decoder_out.shape}")

        # input_size = tf.shape(x)[1:3]

        # decoder_out = tf.image.resize(decoder_out, input_size)

        return out

class Deeplabv3(tf.keras.Sequential):

    """ Tensorflow 2 Implementation of """

    def __init__(self,

                 num_classes,

                 kernel_size_initial_conv,

                 num_channels_atrous,

                 num_channels_DCNN=[256, 512, 1024],

                 num_channels_ASPP=256,

                 kernel_size_atrous=3,

                 kernel_size_DCNN=[1, 3],

                 kernel_size_ASPP=[1, 3, 3, 3],

                 padding='same',

                 nonlinearity='relu',

                 use_batchnorm=True,

                 use_bias=True,

                 data_format='channels_last',

                 MultiGrid=[1, 2, 4],

                 rate_ASPP=[1, 6, 12, 18],

                 atrous_output_stride=16,

                 # Not adapted code for any other out stride

                 **kwargs):

        """ Arguments:

            kernel_size_initial_conv: the size of the kernel for the

                                      first convolution

            num_channels_DCNN: touple with the number of channels for the

                               first three blocks of the DCNN

            kernel_size_DCNN: two element touple with the kernel size of the

                              first and last convolution of the resnet_block

                              (First element) and the middle convolution

                              of the resnet_block (Second element)  """

        super(Deeplabv3, self).__init__(**kwargs)

        self.num_classes = num_classes

        self.add(ResNet_Backbone(kernel_size_initial_conv,

                                 num_channels_DCNN,

                                 kernel_size_DCNN,

                                 padding,

                                 nonlinearity,

                                 use_batchnorm,

                                 use_bias,

                                 False,

                                 data_format))

        self.add(Atrous_conv(num_channels_atrous,

                             kernel_size_atrous,

                             MultiGrid,

                             padding,

                             use_batchnorm,

                             'linear',

                             use_bias,

                             data_format,

                             atrous_output_stride))

        self.add(atrous_spatial_pyramid_pooling(num_channels_ASPP,

                                                kernel_size_ASPP,

                                                rate_ASPP,

                                                padding,

                                                use_batchnorm,

                                                'linear',

                                                use_bias,

                                                data_format))

        self.add(aspp_block(1,

                            1,

                            num_classes,

                            padding,

                            use_batchnorm,

                            'linear',

                            use_bias,

                            data_format))

    def call(self, x, training=False):

        out = super(Deeplabv3, self).call(x, training=training)

        if self.num_classes == 1:

            out = tfkl.Activation('sigmoid')(out)

        else:

            out = tfkl.Activation('softmax')(out)

        # Upsample to same size as the input

        input_size = tf.shape(x)[1:3]

        out = tf.image.resize(out, input_size)

        return out

class ResNet_Backbone(tf.keras.Model):

    def __init__(self,

                 kernel_size_initial_conv,

                 num_channels=[256, 512, 1024],

                 kernel_size_blocks=[1, 3],

                 padding='same',

                 nonlinearity='relu',

                 use_batchnorm=True,

                 use_bias=True,

                 use_pooling=False,

                 data_format='channels_last',

                 **kwargs):

        super(ResNet_Backbone, self).__init__(**kwargs)

        self.first_conv = tfkl.Conv2D(num_channels[0],

                                      kernel_size_initial_conv,

                                      strides=2,

                                      padding=padding,

                                      use_bias=use_bias,

                                      data_format=data_format)

        self.max_pool = tfkl.MaxPool2D(pool_size=(2, 2),

                                       padding='valid')

        self.use_pooling = use_pooling

        self.block1 = resnet_block(False,

                                   num_channels[0],

                                   kernel_size_blocks,

                                   padding,

                                   nonlinearity,

                                   use_batchnorm,

                                   use_bias,

                                   data_format)

        self.block2 = resnet_block(True,

                                   num_channels[1],

                                   kernel_size_blocks,

                                   padding,

                                   nonlinearity,

                                   use_batchnorm,

                                   use_bias,

                                   data_format)

        self.block3 = resnet_block(True,

                                   num_channels[2],

                                   kernel_size_blocks,

                                   padding,

                                   nonlinearity,

                                   use_batchnorm,

                                   use_bias,

                                   data_format)

    def call(self, x, training=False):

        x = self.first_conv(x, training=training)  # output stride 2

        if self.use_pooling:

            x = self.max_pool(x)  # output stride 4

        x = self.block1(x, training=training)  # output stride 2 or 4

        x = self.block2(x, training=training)  # output stride 4 or 8

        x = self.block3(x, training=training)  # output stride 8 or 16

        return x

# full pre-activation residual unit

class resnet_block(tf.keras.Model):

    def __init__(self,

                 use_stride,

                 num_channels,

                 kernel_size=[1, 3],

                 padding='same',

                 nonlinearity='relu',

                 use_batchnorm=True,

                 use_bias=True,

                 data_format='channels_last',

                 **kwargs):

        super(resnet_block, self).__init__(**kwargs)

        self.use_stride = use_stride

        inner_num_channels = num_channels // 4

        if use_stride:

            self.input_conv = basic_conv_block(num_channels,

                                               1,

                                               2,

                                               padding,

                                               nonlinearity,

                                               use_batchnorm,

                                               use_bias,

                                               data_format)

            stride = 2

        else:

            stride = 1

        self.first_conv = basic_conv_block(inner_num_channels,

                                           kernel_size[0],

                                           stride,

                                           padding,

                                           nonlinearity,

                                           use_batchnorm,

                                           use_bias,

                                           data_format)

        self.second_conv = basic_conv_block(inner_num_channels,

                                            kernel_size[1],

                                            1,

                                            padding,

                                            nonlinearity,

                                            use_batchnorm,

                                            use_bias,

                                            data_format)

        self.third_conv = basic_conv_block(num_channels,

                                           kernel_size[0],

                                           1,

                                           padding,

                                           nonlinearity,

                                           use_batchnorm,

                                           use_bias,

                                           data_format)

    def call(self, x, training=False):

        residual = self.first_conv(x, training=training)

        if self.use_stride:

            x = self.input_conv(x, training=training)

        residual = self.second_conv(residual, training=training)

        residual = self.third_conv(residual, training=training)

        output = tfkl.Add()([residual, x])

        return output

class basic_conv_block(tf.keras.Sequential):

    def __init__(self,

                 num_channels,

                 kernel_size,

                 stride=1,

                 padding='same',

                 nonlinearity='relu',

                 use_batchnorm=True,

                 use_bias=True,

                 data_format='channels_last',

                 rate=1,

                 **kwargs):

        super(basic_conv_block, self).__init__(**kwargs)

        if use_batchnorm:

            self.add(tfkl.BatchNormalization(axis=-1,

                                             momentum=0.95,

                                             epsilon=0.001))

        self.add(tfkl.Activation(nonlinearity))

        self.add(tfkl.Conv2D(num_channels,

                             kernel_size,

                             strides=stride,

                             padding=padding,

                             use_bias=use_bias,

                             data_format=data_format,

                             dilation_rate=rate))

    def call(self, x, training=False):

        output = super(basic_conv_block, self).call(x, training=training)

        return output

# ####################### Atrous Convolution ####################### #

class Atrous_conv(tf.keras.Model):

    def __init__(self,

                 num_channels,

                 kernel_size=3,

                 MultiGrid=[1, 2, 4],

                 padding='same',

                 use_batchnorm=True,

                 nonlinearity='linear',

                 use_bias=True,

                 data_format='channels_last',

                 output_stride=16,

                 **kwargs):

        super(Atrous_conv, self).__init__(**kwargs)

        if output_stride == 16:

            multiplier = 2

        else:

            multiplier = 1

        self.first_conv = basic_conv_block(num_channels,

                                           kernel_size,

                                           1,

                                           padding,

                                           nonlinearity,

                                           use_batchnorm,

                                           use_bias,

                                           data_format,

                                           rate=int(multiplier * MultiGrid[0]))

        self.second_conv = basic_conv_block(num_channels,

                                            kernel_size,

                                            1,

                                            padding,

                                            nonlinearity,

                                            use_batchnorm,

                                            use_bias,

                                            data_format,

                                            rate=int(multiplier * MultiGrid[1]))

        self.third_conv = basic_conv_block(num_channels,

                                           kernel_size,

                                           1,

                                           padding,

                                           nonlinearity,

                                           use_batchnorm,

                                           use_bias,

                                           data_format,

                                           rate=int(multiplier * MultiGrid[2]))

    def call(self, x, training=False):

        x = self.first_conv(x, training)

        x = self.second_conv(x, training)

        x = self.third_conv(x, training)

        return x

# ####################### ASPP ####################### #

class atrous_spatial_pyramid_pooling(tf.keras.Model):

    def __init__(self,

                 num_channels=256,

                 kernel_size=[1, 3, 3, 3],

                 rate=[1, 6, 12, 18],

                 padding='same',

                 use_batchnorm=True,

                 nonlinearity='linear',

                 use_bias=True,

                 data_format='channels_last',

                 **kwargs):

        super(atrous_spatial_pyramid_pooling, self).__init__(**kwargs)

        self.block_list = []

        self.basic_conv1 = tfkl.Conv2D(num_channels,

                                       kernel_size=1,

                                       padding=padding)

        self.basic_conv2 = tfkl.Conv2D(num_channels,

                                       kernel_size=1,

                                       padding=padding)

        for i in range(len(kernel_size)):

            self.block_list.append(aspp_block(kernel_size[i],

                                              rate[i],

                                              num_channels,

                                              padding,

                                              use_batchnorm,

                                              nonlinearity,

                                              use_bias,

                                              data_format))

    def call(self, x, training=False):

        feature_map_size = tf.shape(x)

        output_list = []

        # Non diluted convolution

        y = tf.math.reduce_mean(x, axis=[1, 2], keepdims=True)  # ~ Average Pooling

        y = self.basic_conv1(y, training=training)

        output_list.append(tf.image.resize(y, (feature_map_size[1], feature_map_size[2])))  # ~ Upsampling

        # Series of diluted convolutions with rates (1, 6, 12, 18)

        for i, block in enumerate(self.block_list):

            output_list.append(block(x, training=training))

        # concatenate all outputs

        out = tf.concat(output_list, axis=3)

        out = self.basic_conv2(out, training=training)

        return out

class aspp_block(tf.keras.Sequential):

    def __init__(self,

                 kernel_size,

                 rate,

                 num_channels=256,

                 padding='same',

                 use_batchnorm=True,

                 nonlinearity='linear',

                 use_bias=True,

                 data_format='channels_last',

                 **kwargs):

        super(aspp_block, self).__init__(**kwargs)

        self.add(tfkl.Conv2D(num_channels,

                             kernel_size,

                             padding=padding,

                             use_bias=use_bias,

                             data_format=data_format,

                             dilation_rate=rate))

        if use_batchnorm:

            self.add(tfkl.BatchNormalization(axis=-1,

                                             momentum=0.95,

                                             epsilon=0.001))

        self.add(tfkl.Activation(nonlinearity))

    def call(self, x, training=False):

        output = super(aspp_block, self).call(x, training=training)

        return output

# ####################### Decoder ####################### #

class Decoder(tf.keras.Model):

    def __init__(self,

                 num_channels_from_backcone=[48],

                 num_channels_UpConv=[512, 256, 128],

                 kernel_size_UpConv=3,

                 stride_UpConv=(2, 2),

                 use_batchnorm_UpConv=False,

                 use_transpose_UpConv=False,

                 use_bias=True,

                 padding='same',

                 data_format='channels_last',

                 **kwargs):

        super(Decoder, self).__init__(**kwargs)

        self.first_conv1x1 = tfkl.Conv2D(num_channels_from_backcone[0],

                                         kernel_size=1,

                                         padding=padding,

                                         data_format=data_format)

        self.second_conv1x1 = tfkl.Conv2D(num_channels_from_backcone[1],

                                          kernel_size=1,

                                          padding=padding,

                                          data_format=data_format)

        self.conv1 = Up_Conv2D(num_channels_conv=num_channels_UpConv[0],

                               kernel_size=kernel_size_UpConv,

                               use_batchnorm=use_batchnorm_UpConv,

                               use_transpose=use_transpose_UpConv,

                               strides=stride_UpConv)

        self.conv2 = Up_Conv2D(num_channels_conv=num_channels_UpConv[1],

                               kernel_size=kernel_size_UpConv,

                               use_batchnorm=use_batchnorm_UpConv,

                               use_transpose=use_transpose_UpConv,

                               strides=stride_UpConv)

        self.conv3 = Up_Conv2D(num_channels_conv=num_channels_UpConv[2],

                               kernel_size=kernel_size_UpConv,

                               use_batchnorm=use_batchnorm_UpConv,

                               use_transpose=use_transpose_UpConv,

                               strides=stride_UpConv)

    def call(self, in_atrous, in_encoder, in_DCNN, training=False):

        in_atrous = self.first_conv1x1(in_atrous, training=training)

        in_DCNN = self.second_conv1x1(in_DCNN, training=training)

        out = tf.concat([in_atrous, in_encoder], axis=3)

        out = self.conv1(out, training=training)

        out = tf.concat([in_DCNN, out], axis=3)

        out = self.conv2(out, training=training)

        out = self.conv3(out, training=training)

        return out

class Up_Conv2D(tf.keras.Sequential):

    def __init__(self,

                 num_channels_conv,

                 num_channels_UpConv=256,

                 kernel_size=3,

                 nonlinearity='relu',

                 use_batchnorm=False,

                 use_transpose=False,

                 use_bias=True,

                 strides=(2, 2),

                 padding='same',

                 data_format='channels_last',

                 **kwargs):

        super(Up_Conv2D, self).__init__(**kwargs)

        if use_transpose:

            self.add(tfkl.Conv2DTranspose(num_channels_UpConv,

                                          kernel_size,

                                          padding='same',

                                          strides=strides,

                                          data_format=data_format))

        else:

            self.add(tfkl.UpSampling2D(size=strides))

        self.add(aspp_block(kernel_size=kernel_size,

                            rate=1,

                            num_channels=num_channels_conv,

                            padding=padding,

                            use_batchnorm=use_batchnorm,

                            nonlinearity=nonlinearity,

                            use_bias=use_bias,

                            data_format=data_format))

    def call(self, x, training=False):

        out = super(Up_Conv2D, self).call(x, training=training)

        return out