#!/usr/bin/env python

# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

#     http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

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

"""execution script."""

import code

import argparse

import os, warnings

import time

import pandas as pd

import pickle

import sys

import numpy as np

import torch

import utils.exp_utils as utils

from evaluator import Evaluator

from predictor import Predictor

from plotting import plot_batch_prediction

from datetime import datetime

for msg in ["Attempting to set identical bottom==top results",

            "This figure includes Axes that are not compatible with tight_layout",

            "Data has no positive values, and therefore cannot be log-scaled.",

            ".*invalid value encountered in double_scalars.*",

            ".*Mean of empty slice.*"]:

    warnings.filterwarnings("ignore", msg)

def train(logger):

    """

    perform the training routine for a given fold. saves plots and selected parameters to the experiment dir

    specified in the configs.

    """

    time_start_train = time.time()

    logger.info('performing training in {}D over fold {} on experiment {} with model {}'.format(

        cf.dim, cf.fold, cf.exp_dir, cf.model))

    net = model.net(cf, logger).cuda()

    if hasattr(cf, "optimizer") and cf.optimizer.lower() == "adam":

        logger.info("Using Adam optimizer.")

        optimizer = torch.optim.Adam(utils.parse_params_for_optim(net, weight_decay=cf.weight_decay,

                                                                   exclude_from_wd=cf.exclude_from_wd),

                                      lr=cf.learning_rate[0])

    else:

        logger.info("Using AdamW optimizer.")

        optimizer = torch.optim.AdamW(utils.parse_params_for_optim(net, weight_decay=cf.weight_decay,

                                                                   exclude_from_wd=cf.exclude_from_wd),

                                      lr=cf.learning_rate[0])

    if cf.dynamic_lr_scheduling:

        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode=cf.scheduling_mode, factor=cf.lr_decay_factor,

                                                               patience=cf.scheduling_patience)

    model_selector = utils.ModelSelector(cf, logger)

    train_evaluator = Evaluator(cf, logger, mode='train')

    val_evaluator = Evaluator(cf, logger, mode=cf.val_mode)

    starting_epoch = 1

    # prepare monitoring

    monitor_metrics = utils.prepare_monitoring(cf)

    if cf.resume:

        checkpoint_path = os.path.join(cf.fold_dir, "last_checkpoint")

        starting_epoch, net, optimizer, monitor_metrics = \

            utils.load_checkpoint(checkpoint_path, net, optimizer)

        logger.info('resumed from checkpoint {} to epoch {}'.format(checkpoint_path, starting_epoch))

    ####### Use this to create hdf5

    logger.info('loading dataset and initializing batch generators...')

    print ("Starting data_loader.get_train_generators in exec...",datetime.now().strftime("%m/%d/%Y %H:%M:%S:%f"))

    batch_gen = data_loader.get_train_generators(cf, logger)

    print ("Finished data_loader.get_train_generators in exec...",datetime.now().strftime("%m/%d/%Y %H:%M:%S:%f"))

    ####### Writing out train data to file

    #train_data = dict()

    #print ('Write training data to json')

    #for bix in range(cf.num_train_batches):

    #     batch = next(batch_gen['train'])

    #     train_data.update(batch)

    #with open('train_data.json', 'w') as outfile:

    #    json.dump(train_data, outfile)

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

    for epoch in range(starting_epoch, cf.num_epochs + 1):

        logger.info('starting training epoch {}'.format(epoch))

        start_time = time.time()

        net.train()

        train_results_list = []

        for bix in range(cf.num_train_batches):

            ######### Insert call to grab right training data fold from hdf5

            print ("Get next batch_gen['train] ...",datetime.now().strftime("%m/%d/%Y %H:%M:%S:%f"))

            ##Stalled

            batch = next(batch_gen['train']) ######## Instead of this line, grab a batch from training data fold

            tic_fw = time.time()

            print ("Start forward pass...",time.time())

            results_dict = net.train_forward(batch)

            tic_bw = time.time()

            optimizer.zero_grad()

            print ("Start backward pass..",datetime.now().strftime("%m/%d/%Y %H:%M:%S:%f"))

            results_dict['torch_loss'].backward()

            print ("Start optimizing...",datetime.now().strftime("%m/%d/%Y %H:%M:%S:%f"))

            optimizer.step()

            print('\rtr. batch {0}/{1} (ep. {2}) fw {3:.2f}s / bw {4:.2f} s / total {5:.2f} s || '.format(

                bix + 1, cf.num_train_batches, epoch, tic_bw - tic_fw, time.time() - tic_bw,

                time.time() - tic_fw) + results_dict['logger_string'], flush=True, end="")

            print ("Results Dict Size: ",sys.getsizeof(results_dict))

            train_results_list.append(({k:v for k,v in results_dict.items() if k != "seg_preds"}, batch["pid"]))

            print("Loop through train batch DONE",datetime.now().strftime("%m/%d/%Y %H:%M:%S:%f"),(time.time()-time_start_train)/60, "minutes since training started")

        _, monitor_metrics['train'] = train_evaluator.evaluate_predictions(train_results_list, monitor_metrics['train'])

        logger.info('generating training example plot.')

        utils.split_off_process(plot_batch_prediction, batch, results_dict, cf, outfile=os.path.join(

           cf.plot_dir, 'pred_example_{}_train.png'.format(cf.fold)))

        train_time = time.time() - start_time

        logger.info('starting validation in mode {}.'.format(cf.val_mode))

        with torch.no_grad():

            net.eval()

            if cf.do_validation:

                val_results_list = []

                val_predictor = Predictor(cf, net, logger, mode='val')

                for _ in range(batch_gen['n_val']):

                    ########## Insert call to grab right validation data fold from hdf5

                    batch = next(batch_gen[cf.val_mode])

                    if cf.val_mode == 'val_patient':

                        results_dict = val_predictor.predict_patient(batch)

                    elif cf.val_mode == 'val_sampling':

                        results_dict = net.train_forward(batch, is_validation=True)

                    #val_results_list.append([results_dict['boxes'], batch['pid']])

                    val_results_list.append(({k:v for k,v in results_dict.items() if k != "seg_preds"}, batch["pid"]))

                _, monitor_metrics['val'] = val_evaluator.evaluate_predictions(val_results_list, monitor_metrics['val'])

                model_selector.run_model_selection(net, optimizer, monitor_metrics, epoch)

            # update monitoring and prediction plots

            monitor_metrics.update({"lr":

                                        {str(g): group['lr'] for (g, group) in enumerate(optimizer.param_groups)}})

            logger.metrics2tboard(monitor_metrics, global_step=epoch)

            epoch_time = time.time() - start_time

            logger.info('trained epoch {}: took {} ({} train / {} val)'.format(

                epoch, utils.get_formatted_duration(epoch_time, "ms"), utils.get_formatted_duration(train_time, "ms"),

                utils.get_formatted_duration(epoch_time-train_time, "ms")))

            ########### Insert call to grab right validation data fold from hdf5

            batch = next(batch_gen['val_sampling'])

            results_dict = net.train_forward(batch, is_validation=True)

            logger.info('generating validation-sampling example plot.')

            utils.split_off_process(plot_batch_prediction, batch, results_dict, cf, outfile=os.path.join(

                cf.plot_dir, 'pred_example_{}_val.png'.format(cf.fold)))

        # -------------- scheduling -----------------

        if cf.dynamic_lr_scheduling:

            scheduler.step(monitor_metrics["val"][cf.scheduling_criterion][-1])

        else:

            for param_group in optimizer.param_groups:

                param_group['lr'] = cf.learning_rate[epoch-1]

def test(logger):

    """

    perform testing for a given fold (or hold out set). save stats in evaluator.

    """

    test_ix = np.arange((len(os.listdir(cf.pp_test_data_path))/3)-2,dtype=np.int16)

    print ("Test indices size: "+str(len(test_ix))+" and is held out test set "+str(cf.hold_out_test_set))

    logger.info('starting testing model of fold {} in exp {}'.format(cf.fold, cf.exp_dir))

    net = model.net(cf, logger).cuda()

    test_predictor = Predictor(cf, net, logger, mode='test')

    test_evaluator = Evaluator(cf, logger, mode='test')

    ################ Insert call to grab right test data (fold?) from hdf5

    batch_gen = data_loader.get_test_generator(cf, logger)

    ####code.interact(local=locals())

    test_results_list = test_predictor.predict_test_set(batch_gen, return_results=True)

    test_evaluator.evaluate_predictions(test_results_list)

    test_evaluator.score_test_df()

if __name__ == '__main__':

    stime = time.time()

    parser = argparse.ArgumentParser()

    parser.add_argument('-m', '--mode', type=str,  default='train_test',

                        help='one out of: train / test / train_test / analysis / create_exp')

    parser.add_argument('-f','--folds', nargs='+', type=int, default=None,

                        help='None runs over all folds in CV. otherwise specify list of folds.')

    parser.add_argument('--exp_dir', type=str, default='/path/to/experiment/directory',

                        help='path to experiment dir. will be created if non existent.')

    parser.add_argument('--server_env', default=False, action='store_true',

                        help='change IO settings to deploy models on a cluster.')

    parser.add_argument('--data_dest', type=str, default=None, help="path to final data folder if different from config.")

    parser.add_argument('--use_stored_settings', default=False, action='store_true',

                        help='load configs from existing exp_dir instead of source dir. always done for testing, '

                             'but can be set to true to do the same for training. useful in job scheduler environment, '

                             'where source code might change before the job actually runs.')

    parser.add_argument('--resume', action="store_true", default=False,

                        help='if given, resume from checkpoint(s) of the specified folds.')

    parser.add_argument('--exp_source', type=str, default='experiments/toy_exp',

                        help='specifies, from which source experiment to load configs and data_loader.')

    parser.add_argument('--no_benchmark', action='store_true', help="Do not use cudnn.benchmark.")

    parser.add_argument('--cuda_device', type=int, default=0, help="Index of CUDA device to use.")

    parser.add_argument('-d', '--dev', default=False, action='store_true', help="development mode: shorten everything")

    args = parser.parse_args()

    folds = args.folds

    torch.backends.cudnn.benchmark = not args.no_benchmark

    ########### Creating hdf5 

    #if args.mode = 'create_hdf5':

    #    if folds is None:

    #       folds = range(cf.n_cv_splits)

    #    for fold in folds:

    #       create_hdf_foldwise_with_batch_generator_for_train/val/test

    if args.mode == 'train' or args.mode == 'train_test':

        cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, args.use_stored_settings)

        if args.dev:

            folds = [0,1]

            cf.batch_size, cf.num_epochs, cf.min_save_thresh, cf.save_n_models = 3 if cf.dim==2 else 1, 1, 0, 2

            cf.num_train_batches, cf.num_val_batches, cf.max_val_patients = 5, 1, 1

            cf.test_n_epochs =  cf.save_n_models

            cf.max_test_patients = 2

        cf.data_dest = args.data_dest

        logger = utils.get_logger(cf.exp_dir, cf.server_env)

        logger.info("cudnn benchmark: {}, deterministic: {}.".format(torch.backends.cudnn.benchmark,

                                                                     torch.backends.cudnn.deterministic))

        logger.info("sending tensors to CUDA device: {}.".format(torch.cuda.get_device_name(args.cuda_device)))

        data_loader = utils.import_module('dl', os.path.join(args.exp_source, 'data_loader.py'))

        model = utils.import_module('model', cf.model_path)

        logger.info("loaded model from {}".format(cf.model_path))

        if folds is None:

            folds = range(cf.n_cv_splits)

        with torch.cuda.device(args.cuda_device):

            for fold in folds:

                cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))

                cf.fold = fold

                cf.resume = args.resume

                if not os.path.exists(cf.fold_dir):

                    os.mkdir(cf.fold_dir)

                logger.set_logfile(fold=fold)

                train(logger)

                cf.resume = False

                if args.mode == 'train_test':

                    test(logger)

                #Concatenate test results by detection

            if cf.hold_out_test_set == False:

                  test_frames = [pd.read_pickle(os.path.join(cf.test_dir,f)) for f  in os.listdir(cf.test_dir) if '_test_df.pickle' in f]

                  all_preds = pd.concat(test_frames)

                  all_preds.to_csv(os.path.join(cf.test_dir,"all_folds_test.csv"))

                #Concatenate detection raw boxes across folds

                  det_frames = [pd.read_pickle(os.path.join(cf.exp_dir,f,'raw_pred_boxes_list.pickle')) for f in os.listdir(cf.exp_dir) if 'fold_' in f]

                  all_dets=list()

                  for i in det_frames:

                    all_dets.extend(i)

                  with open(os.path.join(cf.exp_dir, 'all_raw_dets.pickle'), 'wb') as handle:

                    pickle.dump(all_dets, handle)

                #Concatenate detection wbc boxes across folds

                  det_frames = [pd.read_pickle(os.path.join(cf.exp_dir,f,'wbc_pred_boxes_list.pickle')) for f in os.listdir(cf.exp_dir) if 'fold_' in f]

                  all_dets=list()

                  for i in det_frames:

                    all_dets.extend(i)

                  with open(os.path.join(cf.exp_dir, 'all_wbc_dets.pickle'), 'wb') as handle:

                    pickle.dump(all_dets, handle)

    elif args.mode == 'test':

        cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, is_training=False, use_stored_settings=True)

        if args.dev:

            folds = [0,1]

            cf.test_n_epochs = 2; cf.max_test_patients = 2

        cf.data_dest = args.data_dest

        logger = utils.get_logger(cf.exp_dir, cf.server_env)

        data_loader = utils.import_module('dl', os.path.join(args.exp_source, 'data_loader.py'))

        model = utils.import_module('model', cf.model_path)

        logger.info("loaded model from {}".format(cf.model_path))

        if folds is None:

            folds = range(cf.n_cv_splits)

        with torch.cuda.device(args.cuda_device):

            for fold in folds:

                cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))

                cf.fold = fold

                logger.set_logfile(fold=fold)

                test(logger)

            if cf.hold_out_test_set == True:

                 test_mode = "_hold_out"

            else:

                test_mode = None

            test_frames = [pd.read_pickle(os.path.join(cf.test_dir,f)) for f  in os.listdir(cf.test_dir) if '_test_df.pickle' in f] 

            all_preds = pd.concat(test_frames)

            all_preds.to_csv(os.path.join(cf.test_dir,"all_folds{}_test.csv".format(test_mode)))

        #Concatenate detection raw boxes across folds

            det_frames = [pd.read_pickle(os.path.join(cf.exp_dir,f,'raw_pred_boxes{}_list.pickle'.format(test_mode))) for f in os.listdir(cf.exp_dir) if 'fold_' in f]

            all_dets=list()

            for i in det_frames:

              all_dets.extend(i)

            with open(os.path.join(cf.exp_dir, 'all_raw{}_dets.pickle'.format(test_mode)), 'wb') as handle:

              pickle.dump(all_dets, handle)

        #Concatenate detection wbc boxes across folds

            det_frames = [pd.read_pickle(os.path.join(cf.exp_dir,f,'wbc_pred_boxes{}_list.pickle'.format(test_mode))) for f in os.listdir(cf.exp_dir) if 'fold_' in f]

            all_dets=list()

            for i in det_frames:

              all_dets.extend(i)

            with open(os.path.join(cf.exp_dir, 'all_wbc{}_dets.pickle'.format(test_mode)), 'wb') as handle:

              pickle.dump(all_dets, handle)

    # load raw predictions saved by predictor during testing, run aggregation algorithms and evaluation.

    elif args.mode == 'analysis':

        cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, is_training=False, use_stored_settings=True)

        logger = utils.get_logger(cf.exp_dir, cf.server_env)

        if args.dev:

            cf.test_n_epochs = 2

        if cf.hold_out_test_set and cf.ensemble_folds:

            # create and save (unevaluated) predictions across all folds

            predictor = Predictor(cf, net=None, logger=logger, mode='analysis')

            results_list = predictor.load_saved_predictions(apply_wbc=True)

            utils.create_csv_output([(res_dict["boxes"], pid) for res_dict, pid in results_list], cf, logger)

            logger.info('starting evaluation...')

            cf.fold = 'overall_hold_out'

            evaluator = Evaluator(cf, logger, mode='test')

            evaluator.evaluate_predictions(results_list)

            evaluator.score_test_df()

        else:

            fold_dirs = sorted([os.path.join(cf.exp_dir, f) for f in os.listdir(cf.exp_dir) if

                         os.path.isdir(os.path.join(cf.exp_dir, f)) and f.startswith("fold")])

            if folds is None:

                folds = range(cf.n_cv_splits)

            for fold in folds:

                cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))

                cf.fold = fold

                logger.set_logfile(fold=fold)

                if cf.fold_dir in fold_dirs:

                    predictor = Predictor(cf, net=None, logger=logger, mode='analysis')

                    results_list = predictor.load_saved_predictions(apply_wbc=True)

                    logger.info('starting evaluation...')

                    evaluator = Evaluator(cf, logger, mode='test')

                    evaluator.evaluate_predictions(results_list)

                    evaluator.score_test_df()

                else:

                    logger.info("Skipping fold {} since no model parameters found.".format(fold))

    # create experiment folder and copy scripts without starting job.

    # useful for cloud deployment where configs might change before job actually runs.

    elif args.mode == 'create_exp':

        cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, use_stored_settings=False)

        logger = utils.get_logger(cf.exp_dir)

        logger.info('created experiment directory at {}'.format(cf.exp_dir))

    else:

        raise RuntimeError('mode specified in args is not implemented...')

    t = utils.get_formatted_duration(time.time() - stime)

    logger.info("{} total runtime: {}".format(os.path.split(__file__)[1], t))

    del logger