#!/usr/bin/env python

import argparse

import json

import os

import sys

from importlib.machinery import SourceFileLoader

from typing import Tuple

import tensorflow as tf

from utils.Evaluation import evaluate, determine_threshold_on_labeled_patients

from utils.default_config_setup import get_config, get_options, get_datasets, Dataset

base_path = os.path.dirname(os.path.abspath(__file__))

def main(args):

    # reset default graph

    tf.reset_default_graph()

    base_path_trainer = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'trainers', f'{args.trainer}.py')

    base_path_network = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models', f'{args.model}.py')

    trainer = getattr(SourceFileLoader(args.trainer, base_path_trainer).load_module(), args.trainer)

    network = getattr(SourceFileLoader(args.model, base_path_network).load_module(), args.model)

    with open(os.path.join(base_path, args.config), 'r') as f:

        json_config = json.load(f)

    dataset = Dataset.BRAINWEB

    options = get_options(batchsize=args.batchsize, learningrate=args.lr, numEpochs=args.numEpochs, zDim=args.zDim, outputWidth=args.outputWidth,

                          outputHeight=args.outputHeight, slices_start=args.slices_start, slices_end=args.slices_end,

                          numMonteCarloSamples=args.numMonteCarloSamples, config=json_config)

    options['data']['dir'] = options["globals"][dataset.value]

    dataset_hc, dataset_pc = get_datasets(options, dataset=dataset)

    config = get_config(

        trainer=trainer,

        options=options,

        optimizer=args.optimizer,

        intermediateResolutions=args.intermediateResolutions,

        dropout_rate=0.2,

        dataset=dataset_hc

    )

    # handle additional Config parameters e.g. for GMVAE

    for arg in vars(args):

        if hasattr(config, arg):

            setattr(config, arg, getattr(args, arg))

    # Create an instance of the model and train it

    model = trainer(tf.Session(), config, network=network)

    # Train it

    model.train(dataset_hc)

    ########################

    #  Evaluate best dice  #

    #########################

    if not args.threshold:

        # if no threshold is given but a dataset => Best dice evaluation on specific dataset

        if args.ds:

            # evaluate specific dataset

            evaluate_optimal(model, options, args.ds)

            return

        else:

            # evaluate all datasets for best dice without hyper intensity prior

            options['applyHyperIntensityPrior'] = False

            evaluate_optimal(model, options, Dataset.Brainweb)

            evaluate_optimal(model, options, Dataset.MSLUB)

            evaluate_optimal(model, options, Dataset.MSISBI2015)

            # evaluate all datasets for best dice without hyper intensity prior

            options['applyHyperIntensityPrior'] = True

            evaluate_optimal(model, options, Dataset.Brainweb)

            evaluate_optimal(model, options, Dataset.MSLUB)

            evaluate_optimal(model, options, Dataset.MSISBI2015)

    ###############################################

    #  Evaluate generalization to other datasets  #

    ###############################################

    if args.threshold and args.ds:  # only threshold is invalid

        evaluate_with_threshold(model, options, args.threshold, args.ds)

    else:

        options['applyHyperIntensityPrior'] = False

        datasetBrainweb = get_evaluation_dataset(options, Dataset.Brainweb)

        _bestDiceVAL, _threshVAL = determine_threshold_on_labeled_patients([datasetBrainweb], model, options, description='VAL')

        print(f"Optimal threshold on MS Lesion Validation Set without optimal postprocessing: {_threshVAL} (Dice-Score {_bestDiceVAL})")

        # Re-evaluate with the previously determined threshold

        evaluate_with_threshold(model, options, _threshVAL, Dataset.Brainweb)

        evaluate_with_threshold(model, options, _threshVAL, Dataset.MSLUB)

        evaluate_with_threshold(model, options, _threshVAL, Dataset.MSISBI2015)

def evaluate_with_threshold(model, options, threshold, dataset):

    options['applyHyperIntensityPrior'] = False

    options['threshold'] = threshold

    description = lambda ds: f'{type(ds).__name__}-VALthresh_{options["threshold"]}'

    evaluation_dataset = get_evaluation_dataset(options, dataset)

    evaluate(evaluation_dataset, model, options, description=description(evaluation_dataset), epoch=str(options['train']['numEpochs']))

def evaluate_optimal(model, options, dataset):

    hyper_intensity_prior_str = ''

    if options['applyHyperIntensityPrior']:

        hyper_intensity_prior_str = "_wPrior"

    evaluation_dataset = get_evaluation_dataset(options, dataset)

    epochs = str(options['train']['numEpochs'])

    description = f'{type(evaluation_dataset).__name__}_upperbound_{options["threshold"]}{hyper_intensity_prior_str}'

    # Evaluate

    evaluate(evaluation_dataset, model, options, description=description, epoch=epochs)

def get_evaluation_dataset(options, dataset=Dataset.BRAINWEB):

    options['data']['dir'] = options["globals"][dataset.value]

    return get_datasets(options, dataset=dataset)[1]

if __name__ == '__main__':

    args = argparse.ArgumentParser(description='Framework')

    args.print_help(sys.stderr)

    args.add_argument('-c', '--config', default='config.default.json', type=str, help='config-path')

    args.add_argument('-b', '--batchsize', default=8, type=int, help='batchsize')

    args.add_argument('-l', '--lr', default=0.0001, type=float, help='learning rate')

    args.add_argument('-E', '--numEpochs', default=1000, type=int, help='how many epochs to train')

    args.add_argument('-z', '--zDim', default=128, type=int, help='Latent dimension')

    args.add_argument('-w', '--outputWidth', default=128, type=int, help='Output width')

    args.add_argument('-g', '--outputHeight', default=128, type=int, help='Output height')

    args.add_argument('-o', '--optimizer', default='ADAM', type=str, help='Can be either ADAM, SGD or RMSProp')

    args.add_argument('-i', '--intermediateResolutions', default=(8, 8), type=Tuple[int], help='Spatial Bottleneck resolution')

    args.add_argument('-s', '--slices_start', default=20, type=int, help='slices start')

    args.add_argument('-e', '--slices_end', default=130, type=int, help='slices end')

    args.add_argument('-t', '--trainer', default='AE', type=str, help='Can be every class from trainers directory')

    args.add_argument('-m', '--model', default='autoencoder', type=str, help='Can be every class from models directory')

    args.add_argument('-O', '--threshold', default=None, type=float, help='Use predefined ThreshOld')

    args.add_argument('-d', '--ds', default=None, type=Dataset, help='Only evaluate on given dataset')

    # following arguments are only relevant for specific architectures

    args.add_argument('-n', '--numMonteCarloSamples', default=0, type=int, help='Amount of Monte Carlos Samples during restoration')

    args.add_argument('-G', '--use_gradient_based_restoration', default=False, type=bool, help='only for ceVAE')

    args.add_argument('-K', '--kappa', default=1.0, type=float, help='only for GANs')

    args.add_argument('-M', '--scale', default=10.0, type=float, help='only for GANs')

    args.add_argument('-R', '--rho', default=1.0, type=float, help='only for ConstrainedAAE')

    args.add_argument('-C', '--dim_c', default=9, type=int, help='only for GMVAE')

    args.add_argument('-Z', '--dim_z', default=128, type=int, help='only for GMVAE')

    args.add_argument('-W', '--dim_w', default=1, type=int, help='only for GMVAE')

    args.add_argument('-A', '--c_lambda', default=1, type=int, help='only for GMVAE')

    args.add_argument('-L', '--restore_lr', default=1e-3, type=float, help='only for GMVAE')

    args.add_argument('-S', '--restore_steps', default=150, type=int, help='only for GMVAE')

    args.add_argument('-T', '--tv_lambda', default=-1.0, type=float, help='only for GMVAE')

    main(args.parse_args())