import tensorflow as tf

import tensorflow.keras.layers as tfkl

import inspect

from Segmentation.model.vnet_build_blocks import Conv_ResBlock, Up_ResBlock

class VNet(tf.keras.Model):

    def __init__(self,

                 num_channels,

                 num_classes,

                 use_2d=False,

                 num_conv_layers=2,

                 kernel_size=(3, 3, 3),

                 activation='prelu',

                 use_batchnorm=True,

                 noise=0.0,

                 dropout_rate=0.25,

                 use_spatial_dropout=True,

                 predict_slice=False,

                 slice_format="mean",

                 **kwargs):

        self.params = str(inspect.currentframe().f_locals)

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

        self.noise = noise

        self.predict_slice = predict_slice

        self.slice_format = slice_format

        block_args = {

            'use_2d': use_2d,

            'num_conv_layers': num_conv_layers,

            'kernel_size': kernel_size,

            'activation': activation,

            'use_batchnorm': use_batchnorm,

            'dropout_rate': dropout_rate,

            'use_spatial_dropout': use_spatial_dropout,

        }

        self.contracting_path = []

        for i in range(len(num_channels)):

            output_ch = num_channels[i]

            self.contracting_path.append(Conv_ResBlock(output_ch,

                                                       **block_args,

                                                       **kwargs))

        self.upsampling_path = []

        n = len(num_channels) - 1

        for i in range(n, -1, -1):

            output_ch = num_channels[i]

            self.upsampling_path.append(Up_ResBlock(output_ch,

                                                    **block_args,

                                                    **kwargs))

        # convolution num_channels at the output

        if use_2d:

            self.conv_output = tfkl.Conv2D(filters=num_channels,

                                           kernel_size=kernel_size,

                                           activation=None,

                                           padding='same')

        else:

            self.conv_output = tfkl.Conv3D(filters=num_classes,

                                           kernel_size=kernel_size,

                                           activation=None,

                                           padding='same')

        if activation == 'prelu':

            self.activation = tfkl.PReLU()  # alpha_initializer=tf.keras.initializers.Constant(value=0.25))

        else:

            self.activation = tfkl.Activation(activation)

        if use_2d:

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

                                        kernel_size=(1, 1),

                                        padding='same')

        else:

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

                                        kernel_size=(1, 1, 1),

                                        padding='same')

        self.output_act = tfkl.Activation('sigmoid' if num_classes == 1 else 'softmax')

    def call(self, x, training):

        if self.noise and training:

            x = tfkl.GaussianNoise(self.noise)(x)

        blocks = []

        # encoder blocks

        for _, down in enumerate(self.contracting_path):

            x, x_before = down(x, training=training)

            blocks.append(x_before)

        # decoder blocks

        for j, up in enumerate(self.upsampling_path):

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

        output = self.conv_output(x)

        output = self.activation(output)

        output = self.conv_1x1(output)

        if self.predict_slice:

            if self.slice_format == "mean":

                output = tf.reduce_mean(output, -4)

                output = tf.expand_dims(output, 1)

            if self.slice_format == "sum":

                output = tf.reduce_sum(output, -4)

                output = tf.expand_dims(output, 1)

        output = self.output_act(output)

        return output