import tensorflow as tf
import tensorflow.keras.layers as tfkl

class SegNet (tf.keras.Model):
    """ Tensorflow 2 Implementation of 'SegNet: A Deep Convolutional Encoder-Decoder
    Architecture for Image Segmentation' https://arxiv.org/abs/1611.09326 """

    def __init__(self,
                 num_channels,
                 num_classes,
                 backbone='default',
                 kernel_size=(3, 3),
                 pool_size=(2, 2),
                 nonlinearity='relu',
                 use_batchnorm=True,
                 use_bias=True,
                 use_transpose=False,
                 use_dropout=False,
                 dropout_rate=0.25,
                 use_spatial_dropout=True,
                 data_format='channels_last',
                 **kwargs):

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

        self.num_classes = num_classes
        self.num_channels = num_channels
        self.backbone = backbone
        self.kernel_size = kernel_size
        self.pool_size = pool_size
        self.nonlinearity = nonlinearity
        self.use_batchnorm = use_batchnorm
        self.use_bias = use_bias
        self.use_transpose = use_transpose
        self.use_dropout = use_dropout
        self.dropout_rate = dropout_rate
        self.use_spatial_dropout = use_spatial_dropout
        self.data_format = data_format

        self.conv_list = tf.keras.Sequential()
        for i in range(len(self.num_channels)):
            output_ch = self.num_channels[i]
            if i == 0 or i == 1:
                num_conv = 2
            else:
                num_conv = 3

            self.conv_list.add(SegNet_Conv2D_Block(output_ch,
                                                   num_conv,
                                                   self.kernel_size,
                                                   self.pool_size,
                                                   self.nonlinearity,
                                                   self.use_batchnorm,
                                                   self.use_bias,
                                                   self.use_dropout,
                                                   self.dropout_rate,
                                                   self.use_spatial_dropout,
                                                   self.data_format))

        self.up_conv_list = tf.keras.Sequential()
        n = len(self.num_channels) - 1

        for j in range(n, -1, -1):
            output = self.num_channels[j]
            if j in [n, n - 1, n - 2]:
                num_conv = 3
            else:
                num_conv = 2
            self.up_conv_list.add(segnet_Up_Conv2D_block(output,
                                                         num_conv_layers=num_conv,
                                                         kernel_size=(2, 2),
                                                         upsampling_size=(2, 2),
                                                         nonlinearity=self.nonlinearity,
                                                         use_batchnorm=self.use_batchnorm,
                                                         use_transpose=self.use_transpose,
                                                         use_bias=self.use_bias,
                                                         strides=(2, 2),
                                                         data_format=self.data_format))

        self.conv_1x1 = tfkl.Conv2D(num_classes,
                                    (1, 1),
                                    activation='linear',
                                    padding='same',
                                    data_format=data_format)

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

        encoded = self.conv_list(x, training=training)
        decoded = self.up_conv_list(encoded, training=training)
        output = self.conv_1x1(decoded)

        if self.num_classes == 1:
            output = tfkl.Activation('sigmoid')(output)
        else:
            output = tfkl.Activation('softmax')(output)
        return output


class SegNet_Conv2D_Block(tf.keras.Sequential):

    def __init__(self,
                 num_channels,
                 num_conv_layers=2,
                 kernel_size=(3, 3),
                 pool_size=(2, 2),
                 nonlinearity='relu',
                 use_batchnorm=True,
                 use_bias=True,
                 use_dropout=False,
                 dropout_rate=0.25,
                 use_spatial_dropout=True,
                 data_format='channels_last',
                 **kwargs):

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

        for _ in range(num_conv_layers):
            self.add(tfkl.Conv2D(num_channels,
                                 kernel_size,
                                 padding='same',
                                 use_bias=use_bias,
                                 data_format=data_format))
            if use_batchnorm:
                self.add(tfkl.BatchNormalization(axis=-1,
                                                 momentum=0.95,
                                                 epsilon=0.001))
            self.add(tfkl.Activation(nonlinearity))

        if use_dropout:
            if use_spatial_dropout:
                self.add(tfkl.SpatialDropout2D(rate=dropout_rate))
            else:
                self.add(tfkl.Dropout(rate=dropout_rate))

        self.add(tfkl.MaxPool2D(pool_size))

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

        output = super(SegNet_Conv2D_Block, self).call(x, training=training)
        return output


class segnet_Up_Conv2D_block(tf.keras.Sequential):

    def __init__(self,
                 num_channels,
                 num_conv_layers,
                 kernel_size=(3, 3),
                 upsampling_size=(2, 2),
                 nonlinearity='relu',
                 use_batchnorm=True,
                 use_transpose=False,
                 use_bias=True,
                 strides=(2, 2),
                 data_format='channels_last',
                 **kwargs):

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

        if use_transpose:
            self.add(tfkl.Conv2DTranspose(num_channels,
                                          kernel_size,
                                          padding='same',
                                          strides=strides,
                                          data_format=data_format))
        else:
            self.add(tf.keras.layers.UpSampling2D(size=upsampling_size))

        for _ in range(num_conv_layers):
            self.add(tfkl.Conv2D(num_channels,
                                 kernel_size,
                                 padding='same',
                                 data_format=data_format))
            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(segnet_Up_Conv2D_block, self).call(x, training=training)
        return output