"""

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 os

from glob import glob

import numpy as np

import matplotlib.pyplot as plt

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

TRAIN_PATH = '/home/ubuntu/data/train_pre/'

VAL_PATH = '/home/ubuntu/data/val_pre/'

TEST_PATH = '/home/ubuntu/data/test_pre/'

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():

    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 generator(path,batch_size):

    lung_mask_list = glob(path + 'final_lung_mask_*.npy')

    nodule_mask_list = glob(path + 'final_nodule_mask_*.npy')

    lung_mask_list.sort()

    nodule_mask_list.sort()

    flag = 0

    start = 0

    cnt = 0

    while (1):

        for i in range(len(lung_mask_list)):

            lung_file = lung_mask_list[i]

            nodule_file = nodule_mask_list[i]

            lung = np.load(lung_file)

            nodule = np.load(nodule_file)

            if(flag):

                lung_train = np.concatenate((lung_train,lung[0:batch_supply]),axis=0).reshape([batch_size,1,512,512])

                nodule_train = np.concatenate((nodule_train,nodule[0:batch_supply]),axis=0).reshape([batch_size,1,512,512])

                yield (lung_train / 255.0, nodule_train / 255.0)

                start = batch_supply

                flag = 0

            while(start + batch_size < len(lung)):

                lung_train = lung[start:start+batch_size].reshape([batch_size,1,512,512])

                nodule_train = nodule[start:start+batch_size].reshape([batch_size,1,512,512])

                yield (lung_train / 255.0, nodule_train / 255.0)

                start += batch_size

            if(start + batch_size == len(lung)):

                lung_train = lung[start:start+batch_size].reshape([batch_size,1,512,512])

                nodule_train = nodule[start:start+batch_size].reshape([batch_size,1,512,512])

                yield (lung_train / 255.0, nodule_train / 255.0)

                flag = 0

                start = 0

            else:

                lung_train = lung[start:]

                nodule_train = nodule[start:]

                batch_supply = batch_size - (len(lung)-start)

                start = 0

                flag = 1

def train_generator(batch_size):

    model = get_unet()

    print('model compileover ...')

    model.fit_generator(generator(TRAIN_PATH,batch_size),steps_per_epoch = 4540, epochs = 2, verbose = 1, validation_data=generator(VAL_PATH,batch_size),validation_steps=650)

    data_pred = np.load(TEST_PATH+'final_lung_mask_9.npy').reshape([1007,1,512,512])

    nodule_true = np.load(TEST_PATH+'final_nodule_mask_9.npy').reshape([1007,1,512,512])

    print(model.evaluate(data_pred,nodule_true,batch_size=2,verbose=1))

if __name__ == '__main__':

    batch_size = 2

    train_generator(batch_size)