from __future__ import print_function

import numpy as np

import keras

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

from keras.layers import Dropout

from sklearn.externals import joblib

import argparse

from keras.callbacks import *

import sys

import theano

import theano.tensor as T

from keras import initializations

from keras.layers import BatchNormalization

import copy

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

'''

    DEFAULT CONFIGURATIONS

'''

def get_options():   

    parser = argparse.ArgumentParser(description='UNET for Lung Nodule Detection')

    parser.add_argument('-out_dir', action="store", default='/scratch/cse/dual/cs5130287/Luna2016/output_final/',

                        dest="out_dir", type=str)

    parser.add_argument('-epochs', action="store", default=500, dest="epochs", type=int)

    parser.add_argument('-batch_size', action="store", default=2, dest="batch_size", type=int)    

    parser.add_argument('-lr', action="store", default=0.001, dest="lr", type=float)

    parser.add_argument('-load_weights', action="store", default=False, dest="load_weights", type=bool)

    parser.add_argument('-filter_width', action="store", default=3, dest="filter_width",type=int)

    parser.add_argument('-stride', action="store", default=3, dest="stride",type=int)

    parser.add_argument('-model_file', action="store", default="", dest="model_file",type=str) #TODO

    parser.add_argument('-save_prefix', action="store", default="model_",

                        dest="save_prefix",type=str)

    opts = parser.parse_args(sys.argv[1:])    

    return opts

def dice_coef(y_true,y_pred):

    y_true = K.flatten(y_true)

    y_pred = K.flatten(y_pred)

    smooth = 0.

    intersection = K.sum(y_true*y_pred)

    return (2. * intersection + smooth) / (K.sum(y_true) + K.sum(y_pred) + smooth)

def dice_coef_loss(y_true, y_pred):

    return 1. - dice_coef(y_true, y_pred)

def gaussian_init(shape, name=None, dim_ordering=None):

   return initializations.normal(shape, scale=0.001, name=name, dim_ordering=dim_ordering)

def get_unet_small(options):

    inputs = Input((1, 512, 512))

    conv1 = Convolution2D(32, options.filter_width, options.stride, activation='elu',border_mode='same')(inputs)

    conv1 = Dropout(0.2)(conv1)

    conv1 = Convolution2D(32, options.filter_width, options.stride, activation='elu',border_mode='same', name='conv_1')(conv1)

    pool1 = MaxPooling2D(pool_size=(2, 2), name='pool_1')(conv1)

    pool1 = BatchNormalization()(pool1)

    conv2 = Convolution2D(64, options.filter_width, options.stride, activation='elu',border_mode='same')(pool1)

    conv2 = Dropout(0.2)(conv2)

    conv2 = Convolution2D(64, options.filter_width, options.stride, activation='elu',border_mode='same', name='conv_2')(conv2)

    pool2 = MaxPooling2D(pool_size=(2, 2), name='pool_2')(conv2)

    pool2 = BatchNormalization()(pool2)

    conv3 = Convolution2D(128, options.filter_width, options.stride, activation='elu',border_mode='same')(pool2)

    conv3 = Dropout(0.2)(conv3)

    conv3 = Convolution2D(128, options.filter_width, options.stride, activation='elu',border_mode='same', name='conv_3')(conv3)

    pool3 = MaxPooling2D(pool_size=(2, 2), name='pool_3')(conv3)

    pool3 = BatchNormalization()(pool3)

    conv4 = Convolution2D(256, options.filter_width, options.stride, activation='elu',border_mode='same')(pool3)

    conv4 = Dropout(0.2)(conv4)

    conv4 = Convolution2D(256, options.filter_width, options.stride, activation='elu',border_mode='same', name='conv_4')(conv4)

    conv4 = BatchNormalization()(conv4)

    # pool4 = MaxPooling2D(pool_size=(2, 2), name='pool_4')(conv4)

    # conv5 = Convolution2D(512, options.filter_width, options.stride, activation='elu',border_mode='same')(pool4)

    # conv5 = Dropout(0.2)(conv5)

    # conv5 = Convolution2D(512, options.filter_width, options.stride, activation='elu',border_mode='same', name='conv_5')(conv5)

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

    # conv6 = Convolution2D(256, options.filter_width, options.stride, activation='elu',border_mode='same')(up6)

    # conv6 = Dropout(0.2)(conv6)

    # conv6 = Convolution2D(256, options.filter_width, options.stride, activation='elu',border_mode='same', name='conv_6')(conv6)

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

    conv7 = Convolution2D(128, options.filter_width, options.stride, activation='elu',border_mode='same')(up7)

    conv7 = Dropout(0.2)(conv7)

    conv7 = Convolution2D(128, options.filter_width, options.stride, activation='elu',border_mode='same', name='conv_7')(conv7)

    conv7 = BatchNormalization()(conv7)

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

    conv8 = Convolution2D(64, options.filter_width, options.stride, activation='elu',border_mode='same')(up8)

    conv8 = Dropout(0.2)(conv8)

    conv8 = Convolution2D(64, options.filter_width, options.stride, activation='elu',border_mode='same', name='conv_8')(conv8)

    conv8 = BatchNormalization()(conv8)

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

    conv9 = Convolution2D(32, options.filter_width, options.stride, activation='elu',border_mode='same')(up9)

    conv9 = Dropout(0.2)(conv9)

    conv9 = Convolution2D(32, options.filter_width, options.stride, activation='elu',border_mode='same', name='conv_9')(conv9)

    conv9 = BatchNormalization()(conv9)

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

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

    model.summary()

    model.compile(optimizer=Adam(lr=options.lr, clipvalue=1., clipnorm=1.), loss=dice_coef_loss, metrics=[dice_coef])

    return model

class WeightSave(Callback):

    def __init__(self, options):

        self.options = options

    def on_train_begin(self, logs={}):

        if self.options.load_weights:            

            print('LOADING WEIGHTS FROM : ' + self.options.model_file)

            weights = joblib.load( self.options.model_file )

            self.model.set_weights(weights)

    def on_epoch_end(self, epochs, logs = {}):

        cur_weights = self.model.get_weights()

        joblib.dump(cur_weights, self.options.save_prefix + '_script_on_epoch_' + str(epochs) + '_lr_' + str(self.options.lr) + '_WITH_STRIDES_' + str(self.options.stride) +'_FILTER_WIDTH_' + str(self.options.filter_width) + '.weights')

class Accuracy(Callback):

    def __init__(self,test_data_x,test_data_y):

        self.test_data_x=test_data_x

        self.test_data_y=test_data_y

        test = T.tensor4('test')

        pred = T.tensor4('pred')

        dc = dice_coef(test,pred)

        self.dc = theano.function([test,pred],dc)

    def on_epoch_end(self,epochs, logs = {}):

        predicted = self.model.predict(self.test_data_x)

        print ("Validation : %f"%self.dc(self.test_data_y,predicted))

def train(use_existing):

    print ("Loading the options ....")

    options = get_options()

    print ("epochs: %d"%options.epochs)

    print ("batch_size: %d"%options.batch_size)

    print ("filter_width: %d"%options.filter_width)

    print ("stride: %d"%options.stride)

    print ("learning rate: %f"%options.lr)

    sys.stdout.flush()

    print('-'*30)

    print('Loading and preprocessing train data...')

    print('-'*30)

    imgs_train = np.load(options.out_dir+"trainImages.npy").astype(np.float32)

    imgs_mask_train = np.load(options.out_dir+"trainMasks.npy").astype(np.float32)

    # Renormalizing the masks

    imgs_mask_train[imgs_mask_train > 0.] = 1.0

    # Now the Test Data

    imgs_test = np.load(options.out_dir+"testImages.npy").astype(np.float32)

    imgs_mask_test_true = np.load(options.out_dir+"testMasks.npy").astype(np.float32)

    # Renormalizing the test masks

    imgs_mask_test_true[imgs_mask_test_true > 0] = 1.0    

    print('-'*30)

    print('Creating and compiling model...')

    print('-'*30)

    model = get_unet_small(options)

    weight_save = WeightSave(options)

    accuracy = Accuracy(copy.deepcopy(imgs_test),copy.deepcopy(imgs_mask_test_true))

    print('-'*30)

    print('Fitting model...')

    print('-'*30)

    model.fit(x=imgs_train, y=imgs_mask_train, batch_size=options.batch_size, nb_epoch=options.epochs, verbose=1, shuffle=True

            ,callbacks=[weight_save, accuracy])

              # callbacks = [accuracy])

              # callbacks=[weight_save,accuracy])

    return model

if __name__ == '__main__':

    # print "epochs"

    model = train(False)