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

# Copyright (C) 2020 Vladimir Juras, Ravinder Regatte and Cem M. Deniz

#

# This file is part of 2019_IWOAI_Challenge

#

# 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/>.

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

import tensorflow as tf

import tf_utilities as tfut

import tf_layers as tflay

import models

import sys

import numpy as np

import re

import time

import os

from functools import partial

import h5py

from sklearn.model_selection import StratifiedKFold

from sklearn.metrics import accuracy_score

from sklearn.preprocessing import label_binarize

from keras.utils import to_categorical

from pathlib import Path

tf.app.flags.DEFINE_boolean('restore', False, 'Whether to restore from previous model.')

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

tf.app.flags.DEFINE_integer('feature', 16, 'Number of root features.')

tf.app.flags.DEFINE_string('model', '4atrous248', 'Model name.')

tf.app.flags.DEFINE_boolean('val', True, 'Whether to use validation.')

tf.app.flags.DEFINE_boolean('full_data', True, 'Whether to use full data set.')

tf.app.flags.DEFINE_float('dr', 1.0, 'Learning rate decay rate.')

tf.app.flags.DEFINE_integer('reso', 384, 'Image size.')

tf.app.flags.DEFINE_integer('slices', 160, 'Number Of Slices')

tf.app.flags.DEFINE_string('loss', 'wce', 'Loss name.')

tf.app.flags.DEFINE_integer('epoch', 400, 'Number of epochs.')

tf.app.flags.DEFINE_boolean('staircase', False, 'If True decay the learning rate at discrete intervals.')

tf.app.flags.DEFINE_integer('seed', 1234, 'Graph-level random seed.')

tf.app.flags.DEFINE_float('dropout', 1.0, 'Dropout rate when training.')

tf.app.flags.DEFINE_string('output_path', None, 'Name of output folder.')

tf.app.flags.DEFINE_boolean('resnet', False, 'Whether to use resnet shortcut.')

tf.app.flags.DEFINE_boolean('early_stopping', True, 'early stopping feature')

tf.app.flags.DEFINE_string('folder', './data', 'Data Folder')

tf.app.flags.DEFINE_integer('noImages', -1, 'how many images to train and validate')

tf.app.flags.DEFINE_float('switchAccuracy', 0.88, 'Training accuracy switch to Dice loss')

tf.app.flags.DEFINE_string('info', ' ', 'add some info to run')

FLAGS = tf.app.flags.FLAGS

switchAccuracy = FLAGS.switchAccuracy

num_classes = 7

num_channels = 1

def _get_cost(logits, batch_y, cost_name='dice', add_regularizers=None, class_weights=None):

    flat_logits = tf.reshape(logits, [-1, num_classes])

    flat_labels = tf.reshape(batch_y, [-1, num_classes])

    if cost_name == 'cross_entropy':

        if class_weights is not None:

            weight_map = tf.multiply(flat_labels, class_weights)

            weight_map = tf.reduce_sum(weight_map, axis=1)

            loss_map = tf.nn.softmax_cross_entropy_with_logits(logits=flat_logits,

                                                            labels=flat_labels)

            weighted_loss = tf.multiply(loss_map, weight_map)

            loss = tf.reduce_mean(weighted_loss)

        else:

            loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=flat_logits, labels=flat_labels))

    elif cost_name == 'dice':

        flat_logits = tf.nn.softmax(flat_logits)[:, 1]

        flat_labels = flat_labels[:, 1]

        inse = tf.reduce_sum(flat_logits*flat_labels)

        l = tf.reduce_sum(flat_logits*flat_logits)

        r = tf.reduce_sum(flat_labels*flat_labels)

        dice = 2 *(inse) / (l+r)

        loss = 1.0-tf.clip_by_value(dice,0,1-1e-10)

    elif cost_name == 'dice_multi':

        dice_multi = 0

        n_classes = num_classes

        for index in range(n_classes):

            flat_logits_ = tf.nn.softmax(flat_logits)[:, index]

            flat_labels_ = flat_labels[:, index]

            inse = tf.reduce_sum(flat_logits_*flat_labels_)

            l = tf.reduce_sum(flat_logits_*flat_logits_)

            r = tf.reduce_sum(flat_labels_*flat_labels_)

            dice = 2 *(inse) / (l+r)

            dice = tf.clip_by_value(dice,0,1-1e-10)

            dice_multi += dice

        loss = n_classes*1.0-dice_multi

    elif cost_name == 'dice_multi_noBG':

        dice_multi = 0

        n_classes = num_classes

        for index in range(1,n_classes):

            flat_logits_ = tf.nn.softmax(flat_logits)[:, index]

            flat_labels_ = flat_labels[:, index]

            inse = tf.reduce_sum(flat_logits_*flat_labels_)

            l = tf.reduce_sum(flat_logits_*flat_logits_)

            r = tf.reduce_sum(flat_labels_*flat_labels_)

            dice = 2 *(inse) / (l+r)

            dice = tf.clip_by_value(dice,0,1-1e-10)

            dice_multi += dice

        loss = (n_classes-1)*1.0-dice_multi

    return loss

def _get_acc(logits, batch_y, cost_name='dice', add_regularizers=None, class_weights=None):

    flat_logits = tf.reshape(logits, [-1, num_classes])

    flat_labels = tf.reshape(batch_y, [-1, num_classes])

    correct_prediction = tf.equal(tf.argmax(flat_logits,1), tf.argmax(flat_labels,1))

    correct_prediction = tf.boolean_mask(correct_prediction, tf.equal(flat_labels[:,0],0))

    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

    return accuracy

def _get_optimizer(start_learning_rate=0.0001, global_step=0, decay_steps=25, decay_rate=0.9):

    learning_rate = tf.train.exponential_decay(start_learning_rate,

                                               global_step,

                                               decay_steps,

                                               decay_rate,

                                               staircase=FLAGS.staircase)

    tf.summary.scalar('learning rate', learning_rate)

    optimizer=tf.train.RMSPropOptimizer(learning_rate=learning_rate, decay=0.995)

    return optimizer

def main(argv=None):

    # if no output path is given, create a new folder using flags

    res = 'res' if FLAGS.resnet else 'nores'

    if FLAGS.output_path is None:

        FLAGS.output_path = 'TrainedModels/' + '_'.join([time.strftime('%m%d_%H%M'),

                                    FLAGS.model,'wceSwitch%.2fDice_AccVal'%(switchAccuracy),

                                    res,

                                    FLAGS.loss, 

                                    'no' + str(FLAGS.noImages),

                                    'reso' + str(FLAGS.reso), 

                                    'features' + str(FLAGS.feature),

                                    'lr' + '{:.1e}'.format(FLAGS.lr), 

                                    'dr' + str(FLAGS.dropout)])

    if not os.path.exists(FLAGS.output_path):

        os.makedirs(FLAGS.output_path)

    # save flags into file

    with open(FLAGS.output_path + '/flags.txt', 'a') as f:

        f.write(str(FLAGS.flag_values_dict()))

    # set seeds for tensorflow and numpy

    tf.set_random_seed(FLAGS.seed)

    np.random.seed(FLAGS.seed)

    # placeholders

    batch_x = tf.placeholder(tf.float32, shape=(None, FLAGS.reso, FLAGS.reso, FLAGS.slices, 1), name='batch_x')

    batch_y = tf.placeholder(tf.float32, shape=(None, None, None, None, num_classes))

    keep_prob = tf.placeholder(tf.float32, shape=[], name='keep_prob')

    global_step = tf.placeholder(tf.int32, shape=[])

    class_weights = tf.placeholder(tf.float32, shape=(num_classes))

    # choose the model

    inference_raw = {'4unet': models.inference_unet4, # the original architecture and use 4 layers

                     '4atrous248': partial(models.inference_atrous4, dilation_rates=[2,4,8])}[FLAGS.model]

    inference = partial(inference_raw, resnet=FLAGS.resnet)

    # get score and probability, add to summary

    score = inference(batch_x, features_root=FLAGS.feature, keep_prob=keep_prob, n_class=num_classes)

    logits = tf.nn.softmax(score)

    # get losses

    dice_cost = _get_cost(score, batch_y, cost_name='dice_multi')

    tf.summary.scalar('dice_loss', dice_cost)  

    dice_cost_noBG = _get_cost(score, batch_y, cost_name='dice_multi_noBG')

    tf.summary.scalar('dice_loss noBG', dice_cost_noBG)   

    cross_entropy = _get_cost(score, batch_y, cost_name='cross_entropy')

    tf.summary.scalar('cross_entropy', cross_entropy) 

    weighted_cross_entropy = _get_cost(score, batch_y, cost_name='cross_entropy', class_weights=class_weights)

    tf.summary.scalar('weighted_cross_entropy',  weighted_cross_entropy)     

    if FLAGS.loss == 'wce': # weighted cross entropy

        cost = weighted_cross_entropy

    elif FLAGS.loss == 'dice': # dice

        cost = dice_cost

    elif FLAGS.loss == 'ce': # cross entropy

        cost = cross_entropy

    else:

        cost = dice_cost

    # get accuracy

    accuracy = _get_acc(score, batch_y)

    # set optimizer with learning rate and decay rate

    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

    with tf.control_dependencies(update_ops):

        with tf.name_scope('rms_optimizer'):

            optimizer = _get_optimizer(FLAGS.lr, global_step, decay_rate=FLAGS.dr)

            optimizer_dice = _get_optimizer(FLAGS.lr, global_step, decay_rate=FLAGS.dr)

            grads = optimizer.compute_gradients(cost)

            grads_dice = optimizer_dice.compute_gradients(dice_cost)

            train = optimizer.apply_gradients(grads)

            train_dice = optimizer_dice.apply_gradients(grads_dice)   

    # get merged summaries

    merged = tf.summary.merge_all()

    # get losses & acc for training

    dice_cost_train = tf.placeholder(tf.float32, shape=[])

    dice_loss_train_summary = tf.summary.scalar('dice_loss_train', dice_cost_train)    

    cross_entropy_train = tf.placeholder(tf.float32, shape=[])

    cross_entropy_train_summary = tf.summary.scalar('cross_entropy_train', cross_entropy_train) 

    weighted_cross_entropy_train = tf.placeholder(tf.float32, shape=[])

    weighted_cross_entropy_train_summary = tf.summary.scalar('weighted_cross_entropy_train',  weighted_cross_entropy_train)    

    accuracy_train = tf.placeholder(tf.float32, shape=[])

    accuracy_train_summary = tf.summary.scalar('accuracy_train',  accuracy_train)  

    # get losses & acc for validation

    dice_cost_val = tf.placeholder(tf.float32, shape=[])

    dice_loss_val_summary = tf.summary.scalar('dice_loss_val', dice_cost_val)    

    cross_entropy_val = tf.placeholder(tf.float32, shape=[])

    cross_entropy_val_summary = tf.summary.scalar('cross_entropy_val', cross_entropy_val) 

    weighted_cross_entropy_val = tf.placeholder(tf.float32, shape=[])

    weighted_cross_entropy_val_summary = tf.summary.scalar('weighted_cross_entropy_val',  weighted_cross_entropy_val)   

    accuracy_val = tf.placeholder(tf.float32, shape=[])

    accuracy_val_summary = tf.summary.scalar('accuracy_val',  accuracy_val)  

    # load data

    #read multiple data

    dataFolder = FLAGS.folder + '/train'

    pathNifti = Path(dataFolder)

    X = []  # create an empty list

    for fileList in list(pathNifti.glob('**/*.im')):

        X.append(fileList)

    X = sorted(X)

    y = []  # create an empty list

    for fileList in list(pathNifti.glob('**/*.seg')):

        y.append(fileList)

    y = sorted(y)

    pathNifti = Path(FLAGS.folder + '/valid')

    X_v = []  # create an empty list

    for fileList in list(pathNifti.glob('**/*.im')):

        X_v.append(fileList)

    X_v = sorted(X_v)

    y_v = []  # create an empty list

    for fileList in list(pathNifti.glob('**/*.seg')):

        y_v.append(fileList)

    y_v = sorted(y_v)

    saver = tf.train.Saver(max_to_keep=0)

    # load mri data and segmentation maps for training

    if FLAGS.noImages ==-1:

        noOfFiles = len(X)

    else:

        noOfFiles = FLAGS.noImages

    list_X = list( X[i] for i in range(noOfFiles) )

    list_y = list( y[i] for i in range(noOfFiles) )

    X_train, y_train, train_info = tfut.loadData_list_h5(list_X,list_y,num_channels)

    print('Dataload is done')

    X_train = tfut.zeroMeanUnitVariance(X_train)

    weights_cross_entropy = tfut.compute_weights_multiClass(y_train,num_classes)

    del list_X, list_y

    # load mri data and segmentation maps for validation

    if FLAGS.noImages ==-1:

        noOfFiles = len(X_v)

    else:

        noOfFiles = FLAGS.noImages

    list_X = list( X_v[i] for i in range(noOfFiles) )

    list_y = list( y_v[i] for i in range(noOfFiles) )

    X_val, y_val, val_info = tfut.loadData_list_h5(list_X, list_y,num_channels)

    X_val = tfut.zeroMeanUnitVariance(X_val)

    del list_X, list_y

    X_train = X_train[...,np.newaxis]

    X_val = X_val[...,np.newaxis]

    # # resize data

    if FLAGS.reso != 384:

        input_size= X_train.shape[2]

        X_train = tfut.batch_resize(X_train, input_size=input_size, output_size=FLAGS.reso, order=3)

        y_train = tfut.batch_resize(y_train, input_size=input_size, output_size=FLAGS.reso, order=0)

        X_val = tfut.batch_resize(X_val, input_size=input_size, output_size=FLAGS.reso, order=3)

        y_val = tfut.batch_resize(y_val, input_size=input_size, output_size=FLAGS.reso, order=0)

    sample_size = X_train.shape[0]

    val_size = X_val.shape[0]

    # initialization for early stopping

    if FLAGS.early_stopping:

        best_acc = 0

        wait = 0

        patience = 500

        switchFlag = 1

    config = tf.ConfigProto()

    config.log_device_placement=False

    config.allow_soft_placement =True

    from tensorflow.python.client import device_lib

    with tf.Session(config=config) as sess:

        sess.run(tf.global_variables_initializer())

        modelNo = 0

        if FLAGS.restore:

            ckpt = tf.train.get_checkpoint_state(FLAGS.output_path)

            model_path = ckpt.model_checkpoint_path

            saver.restore(sess, tf.train.latest_checkpoint(FLAGS.output_path))

            print('Model restored from file: %s' % model_path)

            tmp=re.findall('\d+', model_path)

            modelNo = int(tmp[-1])

        train_writer = tf.summary.FileWriter(FLAGS.output_path, sess.graph)

        start = time.clock()

        prediction = sess.run(score, feed_dict={batch_x: X_train[0:1], 

                                        batch_y: y_train[0:1],

                                        global_step:0,

                                        keep_prob:FLAGS.dropout,

                                        class_weights:weights_cross_entropy})

        pred_shape = prediction.shape

        offset0 = (y_train.shape[1] - pred_shape[1]) // 2

        offset1 = (y_train.shape[2] - pred_shape[2]) // 2

        offset2 = (y_train.shape[3] - pred_shape[3]) // 2

        if offset0 == 0 and offset1 == 0 and offset2 == 0:

            print('SAME padding')

        else:

            y_train = y_train[:, offset0:(-offset0), offset1:(-offset1),offset2:(-offset2),:]

            y_val = y_val[:, offset0:(-offset0), offset1:(-offset1),offset2:(-offset2),:]

        for epoch in range(modelNo+1, FLAGS.epoch+1):

            print('train epoch', epoch, 'sample_size', sample_size) 

            # shuffle data at the beginning of every epoch

            shuffled_idx = np.random.permutation(sample_size)

            wce_train, dice_train, ce_train, acc_train = [], [], [], []

            for j in range(sample_size):

                idx = shuffled_idx[j]

                i = (epoch - 1) * sample_size + j + 1

                # Whether to do left-right mirroring

                step = np.random.choice([1,-1]) 

                if switchFlag: 

                    _, loss, dice_loss, cross_entropy_loss, acc = sess.run([train, weighted_cross_entropy, dice_cost, cross_entropy, accuracy], 

                                                                        feed_dict={batch_x: X_train[idx:idx+1, :, :, ::step, :], 

                                                                                    batch_y: y_train[idx:idx+1, :, :, ::step, :],

                                                                                    global_step:epoch-1,

                                                                                    keep_prob:FLAGS.dropout,

                                                                                    class_weights:weights_cross_entropy})

                else:

                     _, loss, dice_loss, cross_entropy_loss, acc = sess.run([train_dice, weighted_cross_entropy, dice_cost, cross_entropy, accuracy], 

                                                                        feed_dict={batch_x: X_train[idx:idx+1, :, :, ::step, :], 

                                                                                    batch_y: y_train[idx:idx+1, :, :, ::step, :],

                                                                                    global_step:epoch-1,

                                                                                    keep_prob:FLAGS.dropout,

                                                                                    class_weights:weights_cross_entropy})

                wce_train.append(loss)

                dice_train.append(dice_loss)

                ce_train.append(cross_entropy_loss)

                acc_train.append(acc)

            # swithc to dice loss when the CE train accuracy is pretty good

            if np.mean(acc_train) > switchAccuracy:

                switchFlag = 0

                print('@@@@ switchtoDicein Epoch#:' ,epoch )

            print('training weighted loss:', np.mean(wce_train), \

                    ', cross entropy loss:', np.mean(ce_train), \

                    ', dice loss:', np.mean(dice_train), \

                    ', accuracy:', np.mean(acc_train))

            summary = sess.run(weighted_cross_entropy_train_summary, feed_dict={weighted_cross_entropy_train:np.mean(wce_train)})

            train_writer.add_summary(summary, epoch)

            summary = sess.run(dice_loss_train_summary, feed_dict={dice_cost_train:np.mean(dice_train)})

            train_writer.add_summary(summary, epoch)

            summary = sess.run(cross_entropy_train_summary, feed_dict={cross_entropy_train:np.mean(ce_train)})

            train_writer.add_summary(summary, epoch)

            summary = sess.run(accuracy_train_summary , feed_dict={accuracy_train:np.mean(acc_train)})

            train_writer.add_summary(summary, epoch)

            if FLAGS.val:

                summary = sess.run(merged, 

                                    feed_dict={batch_x: X_train[:1],

                                                batch_y: y_train[:1],

                                                global_step:epoch-1,

                                                keep_prob:1.0,

                                                class_weights:weights_cross_entropy})

                train_writer.add_summary(summary, epoch)

                wce_val, dice_val, ce_val, acc_val = [], [], [], []

                for j in range(val_size):

                    loss, dice_loss, cross_entropy_loss, acc = sess.run([weighted_cross_entropy, dice_cost, cross_entropy, accuracy], 

                                                                            feed_dict={batch_x: X_val[j:j+1], 

                                                                                        batch_y: y_val[j:j+1],

                                                                                        global_step:epoch-1,

                                                                                        keep_prob:1.0,

                                                                                        class_weights:weights_cross_entropy})

                    wce_val.append(loss)

                    dice_val.append(dice_loss)

                    ce_val.append(cross_entropy_loss)

                    acc_val.append(acc)

                summary = sess.run(weighted_cross_entropy_val_summary, feed_dict={weighted_cross_entropy_val:np.mean(wce_val)})

                train_writer.add_summary(summary, epoch)

                summary = sess.run(dice_loss_val_summary, feed_dict={dice_cost_val:np.mean(dice_val)})

                train_writer.add_summary(summary, epoch)

                summary = sess.run(cross_entropy_val_summary, feed_dict={cross_entropy_val:np.mean(ce_val)})

                train_writer.add_summary(summary, epoch)

                summary = sess.run(accuracy_val_summary, feed_dict={accuracy_val:np.mean(acc_val)})

                train_writer.add_summary(summary, epoch)  

                print('validation weighted loss:', np.mean(wce_val), \

                    ', cross entropy loss:', np.mean(ce_val), \

                    ', dice loss:', np.mean(dice_val), \

                    ', accuracy:', np.mean(acc_val))

                acc = np.mean(acc_val)

                if  acc - 1e-18 > best_acc:

                    best_acc, wait = acc, 0

                    saver.save(sess, FLAGS.output_path+'/model')

                    with open(FLAGS.output_path + '/SavedEpochNo.txt', 'w') as f:

                        f.write(str(epoch))

                else:

                    saver.save(sess, FLAGS.output_path+'/model_lastEpoch')

                    with open(FLAGS.output_path + '/SavedEpochNoLastEpoch.txt', 'w') as f:

                        f.write(str(epoch))

                    wait += 1

                    if wait > patience:

                        print("!!!!Early Stopping on EPOCH %d!!!!" % epoch)

                        break

                print("!!!!BEST: %f, wait %d !!!"%(best_acc, wait))

if __name__ == '__main__':

    tf.app.run()