# -*- coding: utf-8 -*-

from __future__ import print_function, division, absolute_import, unicode_literals

from core import util

from core import ACNN_pos as ACNN

from core import image_util_pos as image_util

from core import unet_pos as unet

import numpy as np

import click

import os

import logging

from datetime import datetime

t = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

os.environ['CUDA_VISIBLE_DEVICES'] = '2'

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'])

@click.command(context_settings=CONTEXT_SETTINGS)

@click.option('--saliency', default='SRSCN', help='Key word for this running time')

@click.option('--run_times', default=1, type=click.IntRange(min=1, clamp=True), help='network training times')

@click.option('--time', default=t, help='the current time or the time when the model to restore was trained')

@click.option('--trainer_learning_rate', default=0.001, type=click.FloatRange(min=1e-8, clamp=True),

              help='network learning rate')

@click.option('--train_validation_batch_size', default=5, type=click.IntRange(min=1),

              help='the number of validation cases')

@click.option('--test_n_files', default=15, type=click.IntRange(min=1), help='the number of test cases')

@click.option('--train_original_search_path', default='../dataset/train_original/*.nii.gz',

              help='search pattern to find all original training data and label images')

@click.option('--train_search_path', default='../dataset/train_data_2d/*.png',

              help='search pattern to find all training data and label images')

@click.option('--train_data_suffix', default='_img.png', help='suffix pattern for the training data images')

@click.option('--train_label_suffix', default='_lab.png', help='suffix pattern for the training label images')

@click.option('--train_shuffle_data', default=True, type=bool,

              help='whether the order of training files should be randomized after each epoch')

@click.option('--train_crop_patch', default=True, type=bool,

              help='whether patches of a certain size need to be cropped for training')

@click.option('--train_patch_size', default=(-1, -1, -1),

              type=(click.IntRange(min=-1), click.IntRange(min=-1), click.IntRange(min=-1)),

              help='size of the training patches')

@click.option('--train_channels', default=1, type=click.IntRange(min=1), help='number of training data channels')

@click.option('--train_n_class', default=4, type=click.IntRange(min=1),

              help='number of training label classes, including the background')

@click.option('--train_contain_foreground', default=False, type=bool,

              help='if the training patches should contain foreground')

@click.option('--train_label_intensity', default=(0, 88, 200, 244), multiple=True,

              type=click.IntRange(min=0), help='list of intensities of the training ground truths')

@click.option('--net_layers', default=5, type=click.IntRange(min=2),

              help='number of convolutional blocks in the down-sampling path')

@click.option('--net_features_root', default=32, type=click.IntRange(min=1),

              help='number of features of the first convolution layer')

@click.option('--net_cost_name', default=u'exponential_logarithmic',

              type=click.Choice(["cross_entropy", "weighted_cross_entropy", "dice_loss",

                                 "generalized_dice_loss", "cross_entropy+dice_loss",

                                 "weighted_cross_entropy+generalized_dice_loss",

                                 "exponential_logarithmic"]), help='type of the cost function')

@click.option('--net_regularizer_type', default='anatomical_constraint_cae',

              type=click.Choice(['L2_norm', 'L1_norm', 'anatomical_constraint_acnn',

                                 'anatomical_constraint_cae']),

              help='type of regularization')

@click.option('--net_regularization_coefficient', default=5e-4, type=click.FloatRange(min=0),

              help='regularization coefficient')

@click.option('--net_acnn_model_path', default='./autoencoder_trained_%s',

              help='path where to restore the ACNN auto-encoder parameters for regularization')

@click.option('--trainer_batch_size', default=8, type=click.IntRange(min=1, clamp=True),

              help='batch size for each training iteration')

@click.option('--trainer_optimizer_name', default='adam', type=click.Choice(['momentum', 'adam']),

              help='type of the optimizer to use (momentum or adam)')

@click.option('--train_model_path', default='./unet_trained_%s_%s/No_%d', help='path where to store checkpoints')

@click.option('--train_training_iters', default=638, type=click.IntRange(min=1),

              help='number of training iterations during each epoch')

@click.option('--train_epochs', default=30, type=click.IntRange(min=1), help='number of epochs')

@click.option('--train_dropout_rate', default=0.2, type=click.FloatRange(min=0, max=1), help='dropout probability')

@click.option('--train_clip_gradient', default=False, type=bool,

              help='whether to apply gradient clipping with L2 norm threshold 1.0')

@click.option('--train_display_step', default=200, type=click.IntRange(min=1),

              help='number of steps till outputting stats')

@click.option('--train_prediction_path', default='./validation_prediction_%s_%s/No_%d',

              help='path where to save predictions on each epoch')

@click.option('--train_restore', default=False, type=bool, help='whether previous model checkpoint need restoring')

@click.option('--test_search_path', default='../dataset/test_data/*.nii.gz',

              help='a search pattern to find all test data and label images')

@click.option('--test_data_suffix', default='_img.nii.gz', help='suffix pattern for the test data images')

@click.option('--test_label_suffix', default='_lab.nii.gz', help='suffix pattern for the test label images')

@click.option('--test_shuffle_data', default=False, type=bool,

              help='whether the order of the loaded test files path should be randomized')

@click.option('--test_channels', default=1, type=click.IntRange(min=1), help='number of test data channels')

@click.option('--test_n_class', default=4, type=click.IntRange(min=1),

              help='number of test label classes, including the background')

@click.option('--test_label_intensity', default=(0, 88, 200, 244), multiple=True,

              type=click.IntRange(min=0),

              help='tuple of intensities of the test ground truths')

@click.option('--test_prediction_path', default=u'./test_prediction_%s_%s/No_%d',

              help='path where to save test predictions')

@click.option('--val_search_path', default='../dataset/val_data/*.nii.gz',

              help='a search pattern to find all validation data and label images')

@click.option('--val_data_suffix', default='_img.nii.gz', help='suffix pattern for the val data images')

@click.option('--val_label_suffix', default='_lab.nii.gz', help='suffix pattern for the val label images')

@click.option('--val_shuffle_data', default=False, type=bool,

              help='whether the order of the loaded val files path should be randomized')

@click.option('--val_channels', default=1, type=click.IntRange(min=1), help='number of val data channels')

@click.option('--val_n_class', default=4, type=click.IntRange(min=1),

              help='number of val label classes, including the background')

@click.option('--val_label_intensity', default=(0, 88, 200, 244), multiple=True,

              type=click.IntRange(min=0),

              help='tuple of intensities of the test ground truths')

@click.option('--train_center_crop', default=True, type=bool,

              help='whether to extract roi from center during training')

@click.option('--train_center_roi', default=(120, 120, 1), multiple=True, type=click.IntRange(min=0),

              help='roi size you want to extract during training')

@click.option('--test_center_crop', default=True, type=bool,

              help='whether to extract roi from center while testing')

@click.option('--test_center_roi', default=(120, 120, 1), multiple=True, type=click.IntRange(min=0),

              help='roi size you want to extract while testing')

@click.option('--pos_parameter', default=5e-4, type=click.FloatRange(min=0),

              help='position loss weight')

def run(run_times, time, train_search_path, train_data_suffix, train_label_suffix, train_shuffle_data, train_crop_patch, train_patch_size, train_channels, train_n_class, train_contain_foreground, train_label_intensity, train_original_search_path,

         net_layers, net_features_root, net_cost_name, net_regularizer_type, net_regularization_coefficient, net_acnn_model_path,

        trainer_batch_size, trainer_optimizer_name, trainer_learning_rate, train_validation_batch_size, train_model_path, train_training_iters, train_epochs, train_dropout_rate, train_clip_gradient, train_display_step, train_prediction_path, train_restore,

        test_search_path, test_data_suffix, test_label_suffix, test_shuffle_data,

        test_channels, test_n_class, test_label_intensity, test_n_files, test_prediction_path,

        val_search_path, val_data_suffix, val_label_suffix, val_shuffle_data,

        val_channels, val_n_class, val_label_intensity, saliency, train_center_crop, train_center_roi, test_center_crop, test_center_roi, pos_parameter

        ):

    if train_restore:

        assert time != t, "The time when the model to restore was trained is not the time now! "

    train_acc_table = np.array([])

    train_dice_table = np.array([])

    train_auc_table = np.array([])

    train_sens_table = np.array([])

    train_spec_table = np.array([])

    test_acc_table = np.array([])

    test_dice_table = np.array([])

    test_auc_table = np.array([])

    for i in range(run_times):

        train_data_provider = image_util.ImageDataProvider(search_path=train_search_path,

                                                           data_suffix=train_data_suffix,

                                                           label_suffix=train_label_suffix,

                                                           shuffle_data=train_shuffle_data,

                                                           crop_patch=train_crop_patch,

                                                           patch_size=train_patch_size,

                                                           channels=train_channels,

                                                           n_class=train_n_class,

                                                           contain_foreground=train_contain_foreground,

                                                           label_intensity=train_label_intensity,

                                                           center_crop=train_center_crop,

                                                           center_roi=train_center_roi,

                                                           inference_phase=False

                                                           )

        train_original_data_provider = image_util.ImageDataProvider(search_path=train_original_search_path,

                                                                    data_suffix=test_data_suffix,

                                                                    label_suffix=test_label_suffix,

                                                                    shuffle_data=False,

                                                                    crop_patch=False,

                                                                    patch_size=train_patch_size,

                                                                    channels=train_channels,

                                                                    n_class=train_n_class,

                                                                    contain_foreground=train_contain_foreground,

                                                                    label_intensity=train_label_intensity,

                                                                    center_crop=train_center_crop,

                                                                    center_roi=train_center_roi,

                                                                    inference_phase=True

                                                           )

        test_data_provider = image_util.ImageDataProvider(search_path=test_search_path,

                                                          data_suffix=test_data_suffix,

                                                          label_suffix=test_label_suffix,

                                                          shuffle_data=test_shuffle_data,

                                                          crop_patch=False,

                                                          channels=test_channels,

                                                          n_class=test_n_class,

                                                          label_intensity=test_label_intensity,

                                                          center_crop=test_center_crop,

                                                          center_roi=test_center_roi,

                                                          inference_phase=True)

        val_data_provider = image_util.ImageDataProvider(search_path=val_search_path,

                                                         data_suffix=val_data_suffix,

                                                         label_suffix=val_label_suffix,

                                                         shuffle_data=val_shuffle_data,

                                                         crop_patch=False,

                                                         channels=val_channels,

                                                         n_class=val_n_class,

                                                         label_intensity=val_label_intensity,

                                                         center_crop=test_center_crop,

                                                         center_roi=test_center_roi,

                                                         inference_phase=True)

        if net_regularizer_type == 'anatomical_constraint_acnn' or net_regularizer_type == 'anatomical_constraint_cae':

            acnn = ACNN.AutoEncoder(batch_size=trainer_batch_size)

            acnn_save_path = acnn.train(train_data_provider, net_acnn_model_path % saliency)

        net = unet.UNet(layers=net_layers, features_root=net_features_root, channels=train_channels,

                        n_class=train_n_class, batch_size=trainer_batch_size, cost_name=net_cost_name,

                        pos_parameter=pos_parameter,

                        cost_kwargs={'regularizer_type': net_regularizer_type,

                                     'regularization_coefficient': net_regularization_coefficient,

                                     'acnn_model_path': (net_acnn_model_path % saliency)})

        trainer = unet.Trainer(net, batch_size=trainer_batch_size, optimizer_name=trainer_optimizer_name,

                               opt_kwargs={'learning_rate': trainer_learning_rate})

        path, train_acc, train_dice, train_auc, train_sens, train_spec = trainer.train(train_data_provider,

                                                                                       val_data_provider,

                                                                                       train_original_data_provider,

                                                                                       train_validation_batch_size,

                                                                                       model_path=train_model_path % (

                                                                                           saliency, time, i),

                                                                                       training_iters=train_training_iters,

                                                                                       epochs=train_epochs,

                                                                                       dropout=train_dropout_rate,

                                                                                       clip_gradient=train_clip_gradient,

                                                                                       display_step=train_display_step,

                                                                                       prediction_path=train_prediction_path % (

                                                                                           saliency, time, i),

                                                                                       restore=train_restore)

        train_acc_table = np.hstack((train_acc_table, train_acc))

        train_dice_table = np.hstack((train_dice_table, train_dice))

        train_auc_table = np.hstack((train_auc_table, train_auc))

        train_sens_table = np.hstack((train_sens_table, train_sens))

        train_spec_table = np.hstack((train_spec_table, train_spec))

        train_summary_path = './train_summary_%s_%s' % (saliency, time)

        if not os.path.exists(train_summary_path):

            logging.info('Allocating {:}'.format(train_summary_path))

            os.makedirs(train_summary_path)

        np.savez(os.path.join(train_summary_path, 'No_%d.npz' % i), acc=train_acc, dice=train_dice, auc=train_auc,

                 sens=train_sens, spec=train_spec)

        test_data_provider.reset_index()

        test_data, test_labels, test_affine, _ = test_data_provider(test_n_files)

        predictions = net.predict(path, test_data)

        test_acc = unet.acc_rate(predictions, test_labels)

        test_dice = unet.dice_score(predictions, test_labels)

        test_auc = unet.auc_score(predictions, test_labels)

        test_acc_table = np.hstack((test_acc_table, test_acc))

        test_dice_table = np.hstack((test_dice_table, test_dice))

        test_auc_table = np.hstack((test_auc_table, test_auc))

        dice_score_path = './dice_score_%s_%s' % (saliency, time)

        if not os.path.exists(dice_score_path):

            logging.info('Allocating {:}'.format(dice_score_path))

            os.makedirs(dice_score_path)

        np.save(os.path.join(dice_score_path, 'No_%d.npy' % i), test_dice)

        print("Mean Dice score= {:.4f}".format(np.mean(test_dice)))

        for j in range(len(test_data)):

            test_data[j] = np.expand_dims(test_data[j], axis=0).transpose((0, 2, 3, 1, 4))

            test_labels[j] = np.expand_dims(test_labels[j], axis=0).transpose((0, 2, 3, 1, 4))

            predictions[j] = np.expand_dims(predictions[j], axis=0).transpose((0, 2, 3, 1, 4))

        util.save_prediction(test_data, test_labels, predictions, test_prediction_path % (saliency, time, i))

        util.save_prediction_1(predictions, test_affine, test_prediction_path % (saliency, time, i))

        util.save_prediction_2(predictions, test_prediction_path % (saliency, time, i))

        test_summary_path = './test_summary_%s_%s' % (saliency, time)

        if not os.path.exists(test_summary_path):

            logging.info('Allocating {:}'.format(test_summary_path))

            os.makedirs(test_summary_path)

        np.savez(os.path.join(test_summary_path, 'No_%d.npz' % i), acc=test_acc, dice=test_dice, auc=test_auc)

    mean_train_acc = np.mean(np.reshape(train_acc_table, [run_times, -1]), axis=0)

    mean_train_dice = np.mean(np.reshape(train_dice_table, [run_times, -1]), axis=0)

    mean_train_auc = np.mean(np.reshape(train_auc_table, [run_times, -1]), axis=0)

    mean_train_sens = np.mean(np.reshape(train_sens_table, [run_times, -1]), axis=0)

    mean_train_spec = np.mean(np.reshape(train_spec_table, [run_times, -1]), axis=0)

    mean_test_acc = np.mean(np.reshape(test_acc_table, [run_times, -1]), axis=0)

    mean_test_dice = np.mean(np.reshape(train_dice_table, [run_times, -1]), axis=0)

    mean_test_auc = np.mean(np.reshape(train_auc_table, [run_times, -1]), axis=0)

    np.savez('./mean_train_summary_%s_%s.npz' % (saliency, time), acc=mean_train_acc, auc=mean_train_auc,

             sens=mean_train_sens, spec=mean_train_spec, dice=mean_train_dice)

    np.savez('./mean_test_summary_%s_%s.npz' % (saliency, time), acc=mean_test_acc, auc=mean_test_auc,

             dice=mean_test_dice)

if __name__ == '__main__':

    run()