# ==============================================================================

# Copyright (C) 2023 Haresh Rengaraj Rajamohan, Tianyu Wang, Kevin Leung, 

# Gregory Chang, Kyunghyun Cho, Richard Kijowski & Cem M. Deniz 

#

# This file is part of OAI-MRI-TKR

#

# This program is free software: you can redistribute it and/or modify

# it under the terms of the GNU Affero General Public License as published

# by the Free Software Foundation, either version 3 of the License, or

# (at your option) any later version.

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License

# along with this program.  If not, see <https://www.gnu.org/licenses/>.

# ==============================================================================

#!/usr/bin/env python3

import h5py

import os.path

import numpy as np

import pandas as pd

import math

import matplotlib

matplotlib.use('Agg')

import matplotlib.pyplot as plt

import tensorflow as tf

#from sklearn.model_selection import StratifiedKFold

from ModelResnet3D import generate_model

from DataGenerator import DataGenerator

#from keras.models import Sequential

#from keras.optimizers import SGD, Adam

#from keras.layers import Dropout, Dense, Conv3D, MaxPooling3D, GlobalAveragePooling3D, Activation, BatchNormalization,Flatten

from keras.callbacks import LearningRateScheduler, TensorBoard, EarlyStopping, ModelCheckpoint, Callback

from sklearn.metrics import roc_auc_score

tf.app.flags.DEFINE_boolean('batch_norm', True, 'Use BN or not')

tf.app.flags.DEFINE_float('lr', 0.0001, 'Initial learning rate.')

tf.app.flags.DEFINE_integer('filters_in_last', 128, 'Number of filters on the last layer')

tf.app.flags.DEFINE_string('file_path', '/gpfs/data/denizlab/Users/hrr288/Radiology_test/', 'Main Folder to Save outputs')

tf.app.flags.DEFINE_integer('val_fold', 1, 'Fold fo cross-validation')

tf.app.flags.DEFINE_string('file_folder','/gpfs/data/denizlab/Datasets/OAI/COR_IW_TSE/', 'Path to HDF5 radiographs of test set')

tf.app.flags.DEFINE_string('csv_path', '/gpfs/data/denizlab/Users/hrr288/Tianyu_dat/COR_TSE/', 'Folder with the fold splits')

FLAGS = tf.app.flags.FLAGS

class roc_callback(Callback):

    def __init__(self,index,val_fold):

        _params = {'dim': (352,352,35),

              'batch_size': 4,

              'n_classes': 2,

              'n_channels': 1,

              'shuffle': False,

              'normalize' : True,

              'randomCrop' : False,

              'randomFlip' : False,

              'flipProbability' : -1,

              'cropDim' : (352,352,35)}

        self.x = DataGenerator(directory = FLAGS.csv_path+'Fold_'+str(val_fold)+'/CV_'+str(index)+'_train.csv', file_folder=FLAGS.file_folder, **_params)

        self.x_val = DataGenerator(directory = FLAGS.csv_path+'Fold_'+str(val_fold)+'/CV_'+str(index)+'_val.csv',file_folder=FLAGS.file_folder,  **_params)

        self.y = pd.read_csv(FLAGS.csv_path+'Fold_'+str(val_fold)+'/CV_'+str(index)+'_train.csv').Label

        self.y_val = pd.read_csv(FLAGS.csv_path+'Fold_'+str(val_fold)+'/CV_'+str(index)+'_val.csv').Label

        self.auc = []

        self.val_auc = []

        self.losses = []

        self.val_losses = []

    def on_train_begin(self, logs={}):

        return

    def on_train_end(self, logs={}):

        return

    def on_epoch_begin(self, epoch, logs={}):

        return

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

        self.losses.append(logs.get('loss'))

        self.val_losses.append(logs.get('val_loss'))

        y_pred = self.model.predict_generator(self.x)

        y_true = self.y[:len(y_pred)]

        roc = roc_auc_score(y_true, y_pred)

        y_pred_val = self.model.predict_generator(self.x_val)

        y_true_val = self.y_val[:len(y_pred_val)]

        roc_val = roc_auc_score(y_true_val, y_pred_val)

        self.auc.append(roc)

        self.val_auc.append(roc_val)

        #print(len(y_true),len(y_true_val))

        print('\rroc-auc: %s - roc-auc_val: %s' % (str(round(roc,4)),str(round(roc_val,4))),end=100*' '+'\n')

        return

    def on_batch_begin(self, batch, logs={}):

        return

    def on_batch_end(self, batch, logs={}):

        return

'''

    Def: Code to plot loss curves

    Params: history = keras output from training

            loss_path = path to save curve

'''

def plot_loss_curves(history, loss_path): #, i):

    f = plt.figure()

    plt.plot(history.history['loss'])

    plt.plot(history.history['val_loss'])

    plt.title('model loss')

    plt.ylabel('loss')

    plt.xlabel('epoch')

    plt.legend(['train', 'validation'], loc='upper left')

    #plt.show()    

    #path = '/data/kl2596/curves/loss/' + loss_path + '.jpeg'

    f.savefig(loss_path)

'''

    Def: Code to plot accuracy curves

    Params: history = keras output from training

            acc_path = path to save curve

'''

def plot_accuracy_curves(history, acc_path): #, i):

    f = plt.figure()

    plt.plot(history.history['acc'])

    plt.plot(history.history['val_acc'])

    plt.title('model accuracy')

    plt.ylabel('accuracy')

    plt.xlabel('epoch')

    plt.legend(['train', 'validation'], loc='upper left')

    #plt.show() 

    #path = '/data/kl2596/curves/accuracy/' + acc_path + '.jpeg'

    f.savefig(acc_path)

def plot_auc_curves(auc_history, acc_path): #, i):

    f = plt.figure()

    plt.plot(auc_history.auc)

    plt.plot(auc_history.val_auc)

    plt.title('model AUC')

    plt.ylabel('auc')

    plt.xlabel('epoch')

    plt.legend(['train', 'validation'], loc='upper left')

    #plt.show() 

    #path = '/data/kl2596/curves/accuracy/' + acc_path + '.jpeg'

    f.savefig(acc_path)

def train_model(model, train_data, val_data, path, index,val_fold):

    #model.summary()

    # Early Stopping callback that can be found on Keras website

    #early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=15)

    # Create path to save weights with model checkpoint

    weights_path = path + 'weights-{epoch:02d}-{val_loss:.2f}-{val_acc:.2f}-{loss:.2f}-{acc:.2f}.hdf5'

    model_checkpoint = ModelCheckpoint(weights_path, monitor = 'val_loss', save_best_only = True, 

                                       verbose=1)

    # Save loss and accuracy curves using Tensorboard

    tensorboard_callback = TensorBoard(log_dir = path, 

                                       histogram_freq = 0, 

                                       write_graph = False, 

                                       write_grads = False, 

                                       write_images = False)

    auc_history = roc_callback(index,val_fold)    

    #callbacks_list = [model_checkpoint, tensorboard_callback, auc_history]

    #es = EarlyStopping(monitor='val_auc', mode='max', verbose=1, patience=50)

    callbacks_list = [model_checkpoint, tensorboard_callback, auc_history]

    history = model.fit_generator(generator = train_data, validation_data = val_data, epochs=10, 

                        #use_multiprocessing=True, workers=6, 

                        callbacks = callbacks_list)

    accuracy = auc_history.val_auc

    print('*****************************')

    print('best auc:',np.max(accuracy))

    print('average auc:',np.mean(accuracy))

    print('*****************************')

    accuracy = history.history['val_acc']

    print('*****************************')

    print('best accuracy:', np.max(accuracy))

    print('average accuracy:', np.mean(accuracy))

    print('*****************************')

    loss_path = path + 'loss_curve.jpeg'

    acc_path = path + 'acc_curve.jpeg'

    auc_path = path + 'auc_curve.jpeg'

    plot_loss_curves(history, loss_path)

    plot_accuracy_curves(history, acc_path)

    plot_auc_curves(auc_history, auc_path)

    #model.save_weights(weights_path)

'''

    Def: Code to run stratified cross validation to train my network

    Params: num_of_folds = number of folds to cross validate

            lr = learning rate

            dr = dropout rate

            filters_in_last = number of filters in last convolutional layer (we tested 64 and 128)

            batch_norm = True or False for batch norm in model

            data = MRI images

            labels = labels corresponding to MRI images

            file_path = path to save network weights, curves, and tensorboard callbacks

'''

def cross_validation(val_fold, lr, filters_in_last, file_path):

    train_params = {'dim': (352,352,35),

          'batch_size': 4,

          'n_classes': 2,

          'n_channels': 1,

          'shuffle': True,

          'normalize' : True,

          'randomCrop' : True,

          'randomFlip' : True,

          'flipProbability' : -1,

          'cropDim' : (352,352,35)}

    val_params = {'dim': (352,352,35),

          'batch_size': 4,

          'n_classes': 2,

          'n_channels': 1,

          'shuffle': False,

          'normalize' : True,

          'randomCrop' : False,

          'randomFlip' : False,

          'flipProbability' : -1,

          'cropDim' : (352,352,35)} 

    model_path = file_path + 'COR_IW_TSE/'

    if not os.path.exists(model_path):

        os.makedirs(model_path)

    #all_folds = [1,2,3,4,5,6,7]

    #train_folds  = all_folds.remove(val_fold)

    num_of_folds = 6

    for i in range(num_of_folds):

        model = generate_model(learning_rate = 2 * 10 **(-4))

        model.summary()

        print(train_params)

        #print(train_index, test_index)

        print('Running Fold', i+1, '/', num_of_folds)   

        fold_path = model_path + 'Fold_' + str(val_fold) + '/CV_'+str(i+1)+'/'

        print(fold_path)

        if not os.path.exists(fold_path):

            os.makedirs(fold_path)    

        training_generator = DataGenerator(directory = FLAGS.csv_path+'Fold_'+str(val_fold)+'/CV_'+str(i+1)+'_train.csv',file_folder=FLAGS.file_folder,  **train_params)

        validation_generator = DataGenerator(directory = FLAGS.csv_path+'Fold_'+str(val_fold)+'/CV_'+str(i+1)+'_val.csv',file_folder=FLAGS.file_folder,  **val_params)

        train_model(model=model, 

                    train_data = training_generator,

                    val_data = validation_generator,

                    path = fold_path, index = i+1,val_fold=val_fold)

def main(argv=None):

    print('Begin training for fold ',FLAGS.val_fold)

    cross_validation(val_fold=FLAGS.val_fold, 

                     lr=FLAGS.lr, filters_in_last=FLAGS.filters_in_last,   

                     file_path = FLAGS.file_path)

if __name__ == "__main__":

    tf.app.run()