import tensorflow as tf

import tensorflow.keras.layers as tfkl

'''The implementation of the 100 layer Tiramisu Network follows

directly from the publication found at https://arxiv.org/pdf/1611.09326.pdf'''

class Hundred_Layer_Tiramisu(tf.keras.Model):

    def __init__(self,

                 growth_rate,

                 layers_per_block,

                 num_channels,

                 num_classes,

                 kernel_size=(3, 3),

                 pool_size=(2, 2),

                 nonlinearity='relu',

                 dropout_rate=0.2,

                 strides=(2, 2),

                 padding='same',

                 use_dropout=False,

                 use_concat=True,

                 **kwargs):

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

        self.growth_rate = growth_rate

        self.layers_per_block = layers_per_block

        self.num_channels = num_channels

        self.num_classes = num_classes

        self.kernel_size = kernel_size

        self.pool_size = pool_size

        self.nonlinearity = nonlinearity

        self.dropout_rate = dropout_rate

        self.strides = strides

        self.padding = padding

        self.use_dropout = use_dropout

        self.use_concat = use_concat

        self.conv_3x3 = tfkl.Conv2D(self.num_channels,

                                    kernel_size,

                                    padding='same')

        self.dense_block_list = []

        self.up_transition_list = []   

        self.conv_1x1 = tfkl.Conv2D(filters=num_classes,

                                    kernel_size=(1, 1),

                                    padding='same')

        layers_counter = 0

        num_filters = num_channels

        print(len(self.layers_per_block))

        for idx in range(0, len(self.layers_per_block)):

            print(idx)

            num_conv_layers = layers_per_block[idx]

            self.dense_block_list.append(dense_layer(num_conv_layers,

                                                     growth_rate,

                                                     kernel_size,

                                                     dropout_rate,

                                                     nonlinearity,

                                                     use_dropout=False,

                                                     use_concat=True))

            layers_counter = layers_counter + num_conv_layers

            num_filters = num_channels + layers_counter * growth_rate

            if idx != len(self.layers_per_block)-1:

                self.dense_block_list.append(down_transition(num_channels=num_filters,

                                                             kernel_size=(1, 1),

                                                             pool_size=(2, 2),

                                                             dropout_rate=0.2,

                                                             nonlinearity='relu',

                                                             use_dropout=False))

        for idx in range(len(self.layers_per_block) - 1, 0, -1):

            num_conv_layers = layers_per_block[idx - 1]

            num_filters = num_conv_layers * growth_rate

            self.up_transition_list.append(up_transition(num_conv_layers,

                                                         num_channels=num_filters,

                                                         growth_rate=self.growth_rate,

                                                         kernel_size=(3, 3),

                                                         strides=(2, 2),

                                                         padding='same',

                                                         use_concat=False))

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

        blocks = []

        x = self.conv_3x3(inputs)

        for i, down in enumerate(self.dense_block_list):

            x = down(x, training=training)

            if i % 2 == 0 and i != len(self.dense_block_list)-1:

                blocks.append(x)

        for i, up in enumerate(self.up_transition_list):

            x = up(x, blocks[- i-1], training=training)

        x = self.conv_1x1(x)

        if self.num_classes == 1:

            output = tfkl.Activation('sigmoid')(x)

        else:

            output = tfkl.Activation('softmax')(x)

        return output

'''------------------------------------------------------------------'''

class conv_layer(tf.keras.Sequential):

    def __init__(self,

                 num_channels,

                 kernel_size=(3, 3),

                 dropout_rate=0.2,

                 nonlinearity='relu',

                 use_dropout=False,

                 **kwargs):

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

        self.num_channels = num_channels

        self.kernel_size = kernel_size

        self.dropout_rate = dropout_rate

        self.nonlinearity = nonlinearity

        self.use_dropout = use_dropout

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

                                         momentum=0.95,

                                         epsilon=0.001))

        self.add(tfkl.Activation(self.nonlinearity))

        self.add(tfkl.Conv2D(self.num_channels,

                             self.kernel_size,

                             padding='same',

                             activation=None, 

                             use_bias=True))

        if use_dropout:

            self.add(tfkl.Dropout(rate=self.dropout_rate))

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

        outputs = super(conv_layer, self).call(inputs, training=training)    

        return outputs

'''-----------------------------------------------------------------'''

class dense_layer(tf.keras.Sequential):

    def __init__(self,

                 num_conv_layers,

                 growth_rate,

                 kernel_size=(3, 3),

                 dropout_rate=0.2,

                 nonlinearity='relu',

                 use_dropout=False,

                 use_concat=True,

                 **kwargs):

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

        self.num_conv_layers = num_conv_layers

        self.growth_rate = growth_rate

        self.kernel_size = kernel_size

        self.dropout_rate = dropout_rate

        self.nonlinearity = nonlinearity

        self.use_dropout = use_dropout

        self.use_concat = use_concat

        self.conv_list = []

        for layer in range(num_conv_layers):

            self.conv_list.append(conv_layer(num_channels=self.growth_rate,

                                             kernel_size=self.kernel_size,

                                             dropout_rate=self.dropout_rate,

                                             nonlinearity=self.nonlinearity,

                                             use_dropout=self.use_dropout))

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

        dense_output = []

        x = inputs

        for i, conv in enumerate(self.conv_list):

            out = conv(x, training=training)

            x = tfkl.concatenate([x, out], axis=-1)

            dense_output.append(out)

        x = tfkl.concatenate(dense_output, axis=-1)

        if self.use_concat:

            x = tfkl.concatenate([x, inputs], axis=-1)

        outputs = x

        return outputs

'''-----------------------------------------------------------------'''

class down_transition(tf.keras.Sequential):

    def __init__(self,

                 num_channels,

                 kernel_size=(1, 1),

                 pool_size=(2, 2),

                 dropout_rate=0.2,

                 nonlinearity='relu',

                 use_dropout=False,

                 **kwargs):

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

        self.kernel_size = kernel_size

        self.pool_size = pool_size

        self.dropout_rate = dropout_rate

        self.nonlinearity = nonlinearity

        self.use_dropout = use_dropout

        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, padding='same'))

        if use_dropout:

            self.add(tfkl.Dropout(rate=self.dropout_rate))

        self.add(tfkl.MaxPooling2D(pool_size))

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

        outputs = super(down_transition, self).call(inputs, training=training)

        return outputs

'''-----------------------------------------------------------------'''

class up_transition(tf.keras.Model):

    def __init__(self,

                 num_conv_layers,

                 num_channels,

                 growth_rate,

                 kernel_size=(3, 3),

                 strides=(2, 2),

                 padding='same',

                 nonlinearity='relu',

                 use_concat=False,

                 **kwargs):

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

        self.num_conv_layers = num_conv_layers

        self.num_channels = num_channels

        self.growth_rate = growth_rate

        self.kernel_size = kernel_size

        self.strides = strides

        self.padding = padding

        self.nonlinearity = nonlinearity

        self.use_concat = use_concat

        self.up_conv = tfkl.Conv2DTranspose(num_channels,

                                            kernel_size,

                                            strides,

                                            padding)

        self.dense_block = dense_layer(num_conv_layers,

                                       growth_rate, 

                                       kernel_size,

                                       strides,

                                       nonlinearity,

                                       use_concat=self.use_concat)

    def call(self, inputs, bridge, training=False):

        up = self.up_conv(inputs, training=training)

        db_up = self.dense_block(up, training=training)

        c_up = tfkl.concatenate([db_up, bridge], axis=3)

        return c_up