##############################################

#                                            #

#                 DO-U-Net                   #

#                   and                      #

#                 DO-SegNet                  #

#                                            #

# Author: Amine Neggazi                      #

# Email: neggazimedlamine@gmail/com          #

# Nick: nemo256                              #

#                                            #

# Please read bc-count/LICENSE               #

#                                            #

##############################################

import tensorflow as tf

import tensorflow_addons as tfa

# custom imports

from config import *

def conv_bn(filters,

            model,

            model_type,

            kernel=(3, 3),

            activation='relu', 

            strides=(1, 1),

            padding='valid',

            type='normal'):

    '''

    This is a custom convolution function:

    :param filters --> number of filters for each convolution

    :param kernel --> the kernel size

    :param activation --> the general activation function (relu)

    :param strides --> number of strides

    :param padding --> model padding (can be valid or same)

    :param type --> to indicate if it is a transpose or normal convolution

    :return --> returns the output after the convolution and batch normalization and activation.

    '''

    if model_type == 'segnet':

        kernel=3

        activation='relu'

        strides=(1, 1)

        padding='same'

        type='normal'

    if type == 'transpose':

        kernel = (2, 2)

        strides = 2

        conv = tf.keras.layers.Conv2DTranspose(filters, kernel, strides, padding)(model)

    else:

        conv = tf.keras.layers.Conv2D(filters, kernel, strides, padding)(model)

    conv = tf.keras.layers.BatchNormalization()(conv)

    conv = tf.keras.layers.Activation(activation)(conv)

    return conv

def max_pool(input):

    '''

    This is a general max pool function with custom parameters.

    '''

    return tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=2)(input)

def concatenate(input1, input2, crop):

    '''

    This is a general concatenation function with custom parameters.

    '''

    return tf.keras.layers.concatenate([tf.keras.layers.Cropping2D(crop)(input1), input2])

def get_callbacks(name):

    '''

    This is a custom function to save only the best checkpoint.

    :param name --> the input model name

    '''

    return [

        tf.keras.callbacks.ModelCheckpoint(f'models/{name}.h5',

                                           save_best_only=True,

                                           save_weights_only=True,

                                           verbose=1)

    ]

# loss functions

@tf.function

def dsc(y_true, y_pred):

    smooth = 1.0

    y_true_f = tf.reshape(y_true, [-1])

    y_pred_f = tf.reshape(y_pred, [-1])

    intersection = tf.reduce_sum(y_true_f * y_pred_f)

    return (2.0 * intersection + smooth) / (tf.reduce_sum(y_true_f) +

                                            tf.reduce_sum(y_pred_f) +

                                            smooth)

@tf.function

def dice_loss(y_true, y_pred):

    return 1 - dsc(y_true, y_pred)

@tf.function

def tversky(y_true, y_pred):

    alpha = 0.7

    smooth = 1.0

    y_true_pos = tf.reshape(y_true, [-1])

    y_pred_pos = tf.reshape(y_pred, [-1])

    true_pos = tf.reduce_sum(y_true_pos * y_pred_pos)

    false_neg = tf.reduce_sum(y_true_pos * (1 - y_pred_pos))

    false_pos = tf.reduce_sum((1 - y_true_pos) * y_pred_pos)

    return (true_pos + smooth) / (true_pos + alpha * false_neg + (1 - alpha) * false_pos + smooth)

@tf.function

def tversky_loss(y_true, y_pred):

    return 1 - tversky(y_true, y_pred)

@tf.function

def focal_tversky(y_true, y_pred):

    return tf.pow((1 - tversky(y_true, y_pred)), 0.75)

@tf.function

def iou(y_true, y_pred):

    intersect = tf.reduce_sum(y_true * y_pred, axis=(1, 2))

    union = tf.reduce_sum(y_true + y_pred, axis=(1, 2))

    return tf.reduce_mean(tf.math.divide_no_nan(intersect, (union - intersect)), axis=1)

@tf.function

def mean_iou(y_true, y_pred):

    y_true_32 = tf.cast(y_true, tf.float32)

    y_pred_32 = tf.cast(y_pred, tf.float32)

    score = tf.map_fn(lambda x: iou(y_true_32, tf.cast(y_pred_32 > x, tf.float32)),

                      tf.range(0.5, 1.0, 0.05, tf.float32),

                      tf.float32)

    return tf.reduce_mean(score)

@tf.function

def iou_loss(y_true, y_pred):

    return -1*mean_iou(y_true, y_pred)

def do_unet():

    '''

    This is the dual output U-Net model.

    It is a custom U-Net with optimized number of layers.

    Please read model.summary()

    '''

    inputs = tf.keras.layers.Input((188, 188, 3))

    # encoder

    filters = 32

    encoder1 = conv_bn(3*filters, inputs, model_type)

    encoder1 = conv_bn(filters, encoder1, model_type, kernel=(1, 1))

    encoder1 = conv_bn(filters, encoder1, model_type)

    pool1 = max_pool(encoder1)

    filters *= 2

    encoder2 = conv_bn(filters, pool1, model_type)

    encoder2 = conv_bn(filters, encoder2, model_type)

    pool2 = max_pool(encoder2)

    filters *= 2

    encoder3 = conv_bn(filters, pool2, model_type)

    encoder3 = conv_bn(filters, encoder3, model_type)

    pool3 = max_pool(encoder3)

    filters *= 2

    encoder4 = conv_bn(filters, pool3, model_type)

    encoder4 = conv_bn(filters, encoder4, model_type)

    # decoder

    filters /= 2

    decoder1 = conv_bn(filters, encoder4, model_type, type='transpose')

    decoder1 = concatenate(encoder3, decoder1, 4)

    decoder1 = conv_bn(filters, decoder1, model_type)

    decoder1 = conv_bn(filters, decoder1, model_type)

    filters /= 2

    decoder2 = conv_bn(filters, decoder1, model_type, type='transpose')

    decoder2 = concatenate(encoder2, decoder2, 16)

    decoder2 = conv_bn(filters, decoder2, model_type)

    decoder2 = conv_bn(filters, decoder2, model_type)

    filters /= 2

    decoder3 = conv_bn(filters, decoder2, model_type, type='transpose')

    decoder3 = concatenate(encoder1, decoder3, 40)

    decoder3 = conv_bn(filters, decoder3, model_type)

    decoder3 = conv_bn(filters, decoder3, model_type)

    out_mask = tf.keras.layers.Conv2D(1, (1, 1), activation='sigmoid', name='mask')(decoder3)

    if cell_type == 'rbc':

        out_edge = tf.keras.layers.Conv2D(1, (1, 1), activation='sigmoid', name='edge')(decoder3)

        model = tf.keras.models.Model(inputs=inputs, outputs=(out_mask, out_edge))

    elif cell_type == 'wbc' or cell_type == 'plt':

        model = tf.keras.models.Model(inputs=inputs, outputs=(out_mask))

    opt = tf.optimizers.Adam(learning_rate=0.0001)

    if cell_type == 'rbc':

        model.compile(loss='mse',

                      loss_weights=[0.1, 0.9],

                      optimizer=opt,

                      metrics=['accuracy'])

    elif cell_type == 'wbc' or cell_type == 'plt':

        model.compile(loss='mse',

                      optimizer=opt,

                      metrics='accuracy')

    return model

def segnet():

    inputs = tf.keras.layers.Input((128, 128, 3))

    # encoder

    filters = 64

    encoder1 = conv_bn(filters, inputs, model_type)

    encoder1 = conv_bn(filters, encoder1, model_type)

    pool1, mask1 = tf.nn.max_pool_with_argmax(encoder1, 3, 2, padding="SAME")

    filters *= 2

    encoder2 = conv_bn(filters, pool1, model_type)

    encoder2 = conv_bn(filters, encoder2, model_type)

    pool2, mask2 = tf.nn.max_pool_with_argmax(encoder2, 3, 2, padding="SAME")

    filters *= 2

    encoder3 = conv_bn(filters, pool2, model_type)

    encoder3 = conv_bn(filters, encoder3, model_type)

    encoder3 = conv_bn(filters, encoder3, model_type)

    pool3, mask3 = tf.nn.max_pool_with_argmax(encoder3, 3, 2, padding="SAME")

    filters *= 2

    encoder4 = conv_bn(filters, pool3, model_type)

    encoder4 = conv_bn(filters, encoder4, model_type)

    encoder4 = conv_bn(filters, encoder4, model_type)

    pool4, mask4 = tf.nn.max_pool_with_argmax(encoder4, 3, 2, padding="SAME")

    encoder5 = conv_bn(filters, pool4, model_type)

    encoder5 = conv_bn(filters, encoder5, model_type)

    encoder5 = conv_bn(filters, encoder5, model_type)

    pool5, mask5 = tf.nn.max_pool_with_argmax(encoder5, 3, 2, padding="SAME")

    # decoder

    unpool1 = tfa.layers.MaxUnpooling2D()(pool5, mask5)

    decoder1 = conv_bn(filters, unpool1, model_type)

    decoder1 = conv_bn(filters, decoder1, model_type)

    decoder1 = conv_bn(filters, decoder1, model_type)

    unpool2 = tfa.layers.MaxUnpooling2D()(decoder1, mask4)

    decoder2 = conv_bn(filters, unpool2, model_type)

    decoder2 = conv_bn(filters, decoder2, model_type)

    decoder2 = conv_bn(filters/2, decoder2, model_type)

    filters /= 2

    unpool3 = tfa.layers.MaxUnpooling2D()(decoder2, mask3)

    decoder3 = conv_bn(filters, unpool3, model_type)

    decoder3 = conv_bn(filters, decoder3, model_type)

    decoder3 = conv_bn(filters/2, decoder3, model_type)

    filters /= 2

    unpool4 = tfa.layers.MaxUnpooling2D()(decoder3, mask2)

    decoder4 = conv_bn(filters, unpool4, model_type)

    decoder4 = conv_bn(filters/2, decoder4, model_type)

    filters /= 2

    unpool5 = tfa.layers.MaxUnpooling2D()(decoder4, mask1)

    decoder5 = conv_bn(filters, unpool5, model_type)

    out_mask = tf.keras.layers.Conv2D(1, (1, 1), activation='sigmoid', name='mask')(decoder5)

    if cell_type == 'rbc':

        out_edge = tf.keras.layers.Conv2D(1, (1, 1), activation='sigmoid', name='edge')(decoder5)

        model = tf.keras.models.Model(inputs=inputs, outputs=(out_mask, out_edge))

    elif cell_type == 'wbc' or cell_type == 'plt':

        model = tf.keras.models.Model(inputs=inputs, outputs=(out_mask))

    opt = tf.optimizers.Adam(learning_rate=0.0001)

    if cell_type == 'rbc':

        model.compile(loss='mse',

                      loss_weights=[0.1, 0.9],

                      optimizer=opt,

                      metrics=[mean_iou, dsc, tversky, 'accuracy'])

    elif cell_type == 'wbc' or cell_type == 'plt':

        model.compile(loss='mse',

                      optimizer=opt,

                      metrics=[mean_iou, dsc, tversky, 'accuracy'])

    return model