# -*- 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()