"""

    Purpose: train a machine learning segmenter that can segment out the nodules on a given 2D patient CT scan slice

    Note:

    - this will train from scratch, with no preloaded weights

    - weights are saved to unet.hdf5 in the specified output folder

"""

from __future__ import print_function

import numpy as np

from keras.models import Model

from keras.layers import Input, merge, Convolution2D, MaxPooling2D, UpSampling2D

from keras.optimizers import Adam

from keras.optimizers import SGD

from keras.callbacks import ModelCheckpoint, LearningRateScheduler

from keras import backend as K

import matplotlib.pyplot as plt

WORKING_PATH = "/home/marshallee/Documents/lung/subset0/"

IMG_ROWS = 512

IMG_COLS = 512

SMOOTH = 1.

K.set_image_dim_ordering('th')  # Theano dimension ordering in this code

def dice_coef(y_true, y_pred):

    y_true_f = K.flatten(y_true)

    y_pred_f = K.flatten(y_pred)

    intersection = K.sum(y_true_f * y_pred_f)

    return (2. * intersection + SMOOTH) / (K.sum(y_true_f) + K.sum(y_pred_f) + SMOOTH)

def dice_coef_loss(y_true, y_pred):

    return -dice_coef(y_true, y_pred)

def dice_coef_np(y_true,y_pred):

    y_true_f = y_true.flatten()

    y_pred_f = y_pred.flatten()

    intersection = np.sum(y_true_f * y_pred_f)

    return (2. * intersection + SMOOTH) / (np.sum(y_true_f) + np.sum(y_pred_f) + SMOOTH)

def get_unet():

    """

        U-net architecture

    """

    inputs = Input((1,IMG_ROWS, IMG_COLS))

    conv1 = Convolution2D(32, 3, 3, activation='relu', border_mode='same')(inputs)

    conv1 = Convolution2D(32, 3, 3, activation='relu', border_mode='same')(conv1)

    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = Convolution2D(64, 3, 3, activation='relu', border_mode='same')(pool1)

    conv2 = Convolution2D(64, 3, 3, activation='relu', border_mode='same')(conv2)

    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

    conv3 = Convolution2D(128, 3, 3, activation='relu', border_mode='same')(pool2)

    conv3 = Convolution2D(128, 3, 3, activation='relu', border_mode='same')(conv3)

    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)

    conv4 = Convolution2D(256, 3, 3, activation='relu', border_mode='same')(pool3)

    conv4 = Convolution2D(256, 3, 3, activation='relu', border_mode='same')(conv4)

    pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)

    conv5 = Convolution2D(512, 3, 3, activation='relu', border_mode='same')(pool4)

    conv5 = Convolution2D(512, 3, 3, activation='relu', border_mode='same')(conv5)

    up6 = merge([UpSampling2D(size=(2, 2))(conv5), conv4], mode='concat', concat_axis=1)

    conv6 = Convolution2D(256, 3, 3, activation='relu', border_mode='same')(up6)

    conv6 = Convolution2D(256, 3, 3, activation='relu', border_mode='same')(conv6)

    up7 = merge([UpSampling2D(size=(2, 2))(conv6), conv3], mode='concat', concat_axis=1)

    conv7 = Convolution2D(128, 3, 3, activation='relu', border_mode='same')(up7)

    conv7 = Convolution2D(128, 3, 3, activation='relu', border_mode='same')(conv7)

    up8 = merge([UpSampling2D(size=(2, 2))(conv7), conv2], mode='concat', concat_axis=1)

    conv8 = Convolution2D(64, 3, 3, activation='relu', border_mode='same')(up8)

    conv8 = Convolution2D(64, 3, 3, activation='relu', border_mode='same')(conv8)

    up9 = merge([UpSampling2D(size=(2, 2))(conv8), conv1], mode='concat', concat_axis=1)

    conv9 = Convolution2D(32, 3, 3, activation='relu', border_mode='same')(up9)

    conv9 = Convolution2D(32, 3, 3, activation='relu', border_mode='same')(conv9)

    conv10 = Convolution2D(1, 1, 1, activation='sigmoid')(conv9)

    model = Model(input=inputs, output=conv10)

    model.compile(optimizer=Adam(lr=1.0e-5), loss=dice_coef_loss, metrics=[dice_coef])

    return model

def test():

    imgs_test = np.load(WORKING_PATH+"trainImages.npy").astype(np.float32)

    imgs_mask_test_true = np.load(WORKING_PATH+"trainMasks.npy").astype(np.float32)

    num = len(imgs_test)

    mean_test = np.mean(imgs_test)  # mean for data centering

    std_test = np.std(imgs_test)  # std for data normalization

    imgs_test -= mean_test  # images should already be standardized, but just in case

    imgs_test /= std_test

    model = get_unet()

    model.load_weights('unet2.hdf5')  

    predMask = model.predict(imgs_test)

    print('pred shape: ', predMask.shape)

    fig,ax = plt.subplots(2,2,figsize=[8,8])

    ax[0,0].imshow(imgs_test[i][0],cmap='gray')

    ax[0,1].imshow(predMask[i][0],cmap='gray')

    ax[1,0].imshow(imgs_test[i][0]*predMask[i][0],cmap='gray')

    plt.show()

if __name__ == '__main__':

    test()