from tensorflow import keras
from layers import *
from loss import *
from tensorflow.keras.layers import MaxPool3D

def unet(pretrained_weights = None,input_size = (64,64,32,1),l_rate=0.01,loss_func=dice_coef_loss):
    inputs = keras.layers.Input(input_size)
    conv1, pool1 = down_sampling(inputs, 8)
    conv2, pool2 = down_sampling(pool1, 16)
    conv3, pool3 = down_sampling(pool2, 32)
    conv4, pool4 = down_sampling(pool3, 64)
    conv5 = keras.layers.Conv3D(128, 3, padding = "same", activation = 'relu')(pool4)
    conv5 = keras.layers.Conv3D(128, 3, padding = "same", activation = 'relu')(conv5)
    up1 = up_sampling(conv5, conv4, 64)
    up2 = up_sampling(up1, conv3, 32)
    up3 = up_sampling(up2, conv2, 16)
    up4 = up_sampling(up3, conv1, 8)

    conv_last = keras.layers.Conv3D(2, 3, padding="same", activation="relu", kernel_initializer = 'he_normal')(up4)
    outputs = keras.layers.Conv3D(1, 1, padding="same", activation="sigmoid")(conv_last)

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

    model.compile(optimizer = tf.keras.optimizers.Adam(learning_rate=l_rate), loss = loss_func, metrics =['accuracy'])
    
    if(pretrained_weights):
    	model.load_weights(pretrained_weights)

    return model

def unet_norm(pretrained_weights = None,input_size = (64,64,32,1),l_rate=0.01,loss_func=dice_coef_loss):
    inputs = keras.layers.Input(input_size)
    conv1, pool1 = down_sampling_with_norm(inputs, 8)
    conv2, pool2 = down_sampling_with_norm(pool1, 16)
    conv3, pool3 = down_sampling_with_norm(pool2, 32)
    conv4, pool4 = down_sampling_with_norm(pool3, 64)
    conv5 = keras.layers.Conv3D(128, 3, padding = "same", activation = 'relu')(pool4)
    conv5 = keras.layers.Conv3D(128, 3, padding = "same", activation = 'relu')(conv5)
    up1 = up_sampling_with_norm(conv5, conv4, 64)
    up2 = up_sampling_with_norm(up1, conv3, 32)
    up3 = up_sampling_with_norm(up2, conv2, 16)
    up4 = up_sampling_with_norm(up3, conv1, 8)

    conv_last = keras.layers.Conv3D(2, 3, padding="same", activation="relu", kernel_initializer = 'he_normal')(up4)
    outputs = keras.layers.Conv3D(1, 1, padding="same", activation="sigmoid")(conv_last)

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

    model.compile(optimizer = tf.keras.optimizers.Adam(learning_rate=l_rate), loss = loss_func, metrics =['Precision','Recall'])
    
    if(pretrained_weights):
    	model.load_weights(pretrained_weights)

    return model

def dense_unet(pretrained_weights = None,input_size = (64,64,32,1),l_rate=0.01,loss_func=dice_coef_loss):
    inputs = keras.layers.Input(input_size)
    conv1, pool1 = down_sampling_with_norm(inputs, 8)
    conv2, pool2 = down_sampling_with_norm(pool1, 16)
    conv3, pool3 = down_sampling_with_norm(pool2, 32)
    dense4 = dense_block(pool3, 64)
    pool4 = MaxPool3D(pool_size=(2, 2, 2))(dense4)
    dense5 = dense_block(pool4, 128)
    
    up1 = up_sampling_with_norm(dense5, dense4, 64)
    up2 = up_sampling_with_norm(up1, conv3, 32)
    up3 = up_sampling_with_norm(up2, conv2, 16)
    up4 = up_sampling_with_norm(up3, conv1, 8)

    conv_last = keras.layers.Conv3D(2, 3, padding="same", activation="relu", kernel_initializer = 'he_normal')(up4)
    outputs = keras.layers.Conv3D(1, 1, padding="same", activation="sigmoid")(conv_last)

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

    model.compile(optimizer = tf.keras.optimizers.Adam(learning_rate=l_rate), loss = loss_func, metrics =['Precision','Recall'])
    
    if(pretrained_weights):
    	model.load_weights(pretrained_weights)

    return model