<reponame>rtgunti/medicaldetectiontoolk
#!/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 argparse
from logging import log
import os
import time
from tokenize import ContStr
from matplotlib.pyplot import box, plot, tight_layout
import torch
import pickle
import numpy as np
import pandas as pd
import copy

import utils.exp_utils as utils
from evaluator import Evaluator
from predictor import Predictor
from plotting import plot_batch_prediction, plot_test_prediction
from torchsummary import summary
import matplotlib.pyplot as plt


def train(logger):
    """
    perform the training routine for a given fold. saves plots and selected parameters to the experiment dir
    specified in the configs.
    """
    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()
    
    optimizer = torch.optim.Adam(net.parameters(), lr=cf.learning_rate[0], weight_decay=cf.weight_decay)
    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, TrainingPlot = utils.prepare_monitoring(cf)
    print("length of monitor_metrics is ", len(monitor_metrics['train']['monitor_values']))
    
    if cf.pre_train_path:
        utils.load_pre_trained_weights(cf.pre_train_path, net, optimizer, cf)
        
    if cf.resume_to_checkpoint:
        path_to_checkpoint = os.path.join(cf.exp_dir, ("fold_" + str(cf.fold)), cf.resume_to_checkpoint)
        starting_epoch, monitor_metrics = utils.load_checkpoint(path_to_checkpoint, net, optimizer, cf)
        logger.info('resumed to checkpoint {} at epoch {}'.format(cf.resume_to_checkpoint, starting_epoch))

    logger.info('loading dataset and initializing batch generators...')
    batch_gen = data_loader.get_train_generators(cf, logger)

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

        logger.info('starting training epoch {}'.format(epoch))
        for param_group in optimizer.param_groups:
            param_group['lr'] = cf.learning_rate[epoch - 1]

        start_time = time.time()

        net.train()
        train_results_list = []

        for bix in range(cf.num_train_batches):
            batch = next(batch_gen['train'])
            results_dict = net.train_forward(batch)
            optimizer.zero_grad()
            results_dict['torch_loss'].backward()
            optimizer.step()
            logger.info('tr. (ep. {0}) batch {1:2}/{2} || '
                        .format(epoch, bix + 1, cf.num_train_batches) + results_dict['logger_string'])
            train_results_list.append([results_dict['boxes'], batch['pid']])
            monitor_metrics['train']['monitor_values'][epoch].append(results_dict['monitor_values'])

        _, monitor_metrics['train'] = train_evaluator.evaluate_predictions(train_results_list, monitor_metrics['train'])
        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']):
                    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)
                        logger.info('val. (ep. {0}) batch {1}/{2} || '
                        .format(epoch, _ + 1, batch_gen['n_val']) + results_dict['logger_string'])                     
                    val_results_list.append([results_dict['boxes'], batch['pid']])
                    monitor_metrics['val']['monitor_values'][epoch].append(results_dict['monitor_values'])

                _, 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
            TrainingPlot.update_and_save(monitor_metrics, starting_epoch, epoch)
            epoch_time = time.time() - start_time
            logger.info('trained epoch {0:}: took {1:.2f} sec. ({2:.2f} train / {3:.2f} val)'.format(
                epoch, epoch_time, train_time, epoch_time-train_time))
            batch = next(batch_gen['val_sampling'])
            results_dict = net.train_forward(batch, is_validation=True)
            logger.info('plotting predictions from validation sampling.')
            plot_batch_prediction(batch, results_dict, cf)

def test(logger):
    """
    perform testing for a given fold (or hold out set). save stats in evaluator.
    """
    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')
    batch_gen = data_loader.get_test_generator(cf, logger)
    test_results_list = test_predictor.predict_test_set(batch_gen, return_results=True)

    # # print("Test without WBC")
    # wbc_bool = False
    # test_results_list = test_predictor.load_saved_predictions(apply_wbc=wbc_bool)
    # test_evaluator.evaluate_predictions(test_results_list)
    # test_evaluator.score_test_df(wbc_flag=wbc_bool)

    # # # print("Test with WBC")
    # wbc_bool = True
    # test_evaluator = Evaluator(cf, logger, mode='test')
    # test_results_list = test_predictor.load_saved_predictions(apply_wbc=wbc_bool)  
    # test_evaluator.evaluate_predictions(test_results_list)
    # test_evaluator.score_test_df(wbc_flag=wbc_bool)

    # get_lesions()
    # load_pred_and_eval()
    # return
    fold = cf.fold
    bool_predict_test_set = False   #set to True to re run the predictions
    bool_predicted = False          #bool used to avoid repeating predictions
    for wbc_bool in [True]:  #[False, True]
        cf.fold = fold
        test_evaluator = Evaluator(cf, logger, mode='test')        
        if bool_predict_test_set and not bool_predicted:
            test_results_list = test_predictor.predict_test_set(batch_gen, return_results=True)
            bool_predicted = True
        test_results_list = test_predictor.load_saved_predictions(apply_wbc=wbc_bool)  
        test_evaluator.evaluate_predictions(test_results_list)
        test_evaluator.score_test_df(wbc_flag=wbc_bool)
        test_results_list = add_det_type_to_test_results(test_results_list, test_evaluator.test_df)
        # character_comparison(test_results_list, test_evaluator.test_df)
        # filter_lesions(test_results_list) 
        # break
        # # continue
        # '''
        # Low confidence lesions = true positives with score < 0.5 (tp with score > 0.5 shown in green)
        # High confidence non-lesions = false positives with score > 0.8 (tp shown in green)
        # '''
        # test_results_list_copy = copy.deepcopy(test_results_list)

        for bool_low_lesion in [False, True]:
            print("wbc_bool : {} bool_low_lesion : {}".format(wbc_bool, bool_low_lesion))
            # plot_test_prediction(copy.deepcopy(test_results_list), cf, wbc_bool, bool_low_lesion)#@rtgunti
            break
        break

    
    # for test_result in test_results_list:
    #     name = test_result[1]
    #     res = test_result[0]
    #     print(name)
    #     for r in res:
    #         print("="*10)
    #         print(len(r))
    #         r = [rr for rr in r if 'box_score' in rr.keys()]
    #         print(len(r))
    #         r = sorted(r, key=lambda k: k['box_score'], reverse=True)
    #         for rr in r[:5]:
    #             print(np.round(rr['box_score'], 6), rr['det_type'])
    # # test_evaluator.evaluate_predictions(test_results_list)
    # print(test_evaluator.test_df.shape)
    # test_evaluator.test_df.sort_values(by=['pred_score'], inplace=True, ascending=False)
    # print(test_evaluator.test_df[test_evaluator.test_df['pid'] == 'd006_int1'].head())

def character_comparison(test_results_list, test_df):
    print(test_df.head())
    for test_result in test_results_list:
        name = test_result[1]
        res = test_result[0]
        print(name)
        for r in res:
            print("="*10)
            print(len(r))
            r = [rr for rr in r if 'box_score' in rr.keys()]
            print(len(r))
            r = sorted(r, key=lambda k: k['box_score'], reverse=True)
            for rr in r[:5]:
                if rr['det_type'] == 'det_tp' and rr['box_score'] < 0.5:
                    rr.update({'char' : 'low_les'})
                elif rr['det_type'] == 'det_fp' and rr['box_score'] > 0.8:
                    rr.update({'char' : 'high_nonles'})
                else:
                    rr.update({'char' : 'det_fp'})
                print(np.round(rr['box_score'], 6), rr['det_type'])


def lcl(logger):
    net = model.net(cf, logger).cuda()
    test_predictor = Predictor(cf, net, logger, mode='test')
    test_evaluator = Evaluator(cf, logger, mode='test')
    batch_gen = data_loader.get_test_generator(cf, logger)

    fold = cf.fold
    cols = ['p_id', 'box_type', 'box_coords', 'box_score', 'box_pred_class_id', 'det_type', 'dims', 'centroid', 'comment']
    df = pd.DataFrame(columns=cols)
    for wbc_bool in [True, False]:
        cf.fold = fold
        wbc_str = 'with_wbc' if wbc_bool else 'without_wbc'
        test_evaluator = Evaluator(cf, logger, mode='test')
        test_results_list = test_predictor.load_saved_predictions(apply_wbc=wbc_bool)          
        test_evaluator.evaluate_predictions(test_results_list)
        test_evaluator.score_test_df(wbc_flag=wbc_bool)
        test_results_list = add_det_type_to_test_results(test_results_list, test_evaluator.test_df)
        # print(test_evaluator.test_df[test_evaluator.test_df['det_type'] == 'det_fn'].head())
        # character_comparison(test_results_list, test_evaluator.test_df)
        les_df = filter_lesions(test_results_list, wbc_str)
        df = df.append(les_df)
    return df

def filter_lesions(test_results_list, wbc_str):
    '''
    filter and print low confidence lesions. test_result_list doesn't include false_negatives
    '''
    cols = ['p_id', 'box_type', 'box_coords', 'box_score', 'box_pred_class_id', 'det_type', 'dims', 'centroid', 'comment']
    lc_lesions = pd.DataFrame(columns=cols)
    hc_non_lesions = pd.DataFrame(columns=cols)
    les_df = pd.DataFrame(columns=cols)
    for test_result_ind, test_result in enumerate(test_results_list):
        name = test_result[1]
        res = test_result[0]
        for r_ind, r in enumerate(res):
            for rr_ind, rr in enumerate(r):
                if 'det_type' in rr and rr['det_type'] == 'det_fn':
                    print(rr['det_type'])
                rr.update({'p_id':name})
                dy, dx, dz = rr['box_coords'][2] - rr['box_coords'][0], rr['box_coords'][3] - rr['box_coords'][1], rr['box_coords'][5] - rr['box_coords'][4]
                rr.update({'dims' : (int(dx), int(dy), int(dz))})
                # rr.update({'dx' : dx, 'dy' : dy, 'dz' : dz})
                rr.update({'centroid' : (int((rr['box_coords'][0] + rr['box_coords'][2])/2.0), int((rr['box_coords'][1] + rr['box_coords'][3])/2.0), int((rr['box_coords'][4] + rr['box_coords'][5])/2.0))})
                if not 'box_score' in rr:
                    ## Try to add fasle_negatives to the test_results_list
                    # print(rr)
                    # rr.update({'comment' : 'low_lesion_' + wbc_str})
                    # rr.update({'box_score' : '-1', 'box_pred_class_id' : -1, 'det_type' : 'na'})
                    # row_df = pd.DataFrame([rr], columns=cols)
                    # lc_lesions = lc_lesions.append(row_df, ignore_index=True)
                    # les_df = les_df.append(row_df, ignore_index=True)
                    continue
                if rr['box_score'] < 0.5 and rr['det_type'] == 'det_tp':
                    rr.update({'comment' : 'low_lesion_' + wbc_str})
                    row_df = pd.DataFrame([rr], columns=cols)
                    lc_lesions = lc_lesions.append(row_df, ignore_index=True)
                    les_df = les_df.append(row_df, ignore_index=True)
                if rr['box_score'] >= 0.5 and rr['det_type'] == 'det_tp':
                    rr.update({'comment' : 'good_tp_' + wbc_str})
                    row_df = pd.DataFrame([rr], columns=cols)
                    les_df = les_df.append(row_df, ignore_index=True)
                if rr['box_score'] > 0.8 and rr['det_type'] == 'det_fp':
                    rr.update({'comment' : 'high_nonlesion_' + wbc_str})
                    row_df = pd.DataFrame([rr], columns=cols)   
                    hc_non_lesions = hc_non_lesions.append(row_df, ignore_index=True) 
                    les_df = les_df.append(row_df, ignore_index=True)   

    print(lc_lesions.shape, hc_non_lesions.shape, les_df.shape)    
    return les_df

def add_det_type_to_test_results(test_results_list, test_df):

    det_cnt = 0
    for test_result_ind, test_result in enumerate(test_results_list):
        name = test_result[1]
        res = test_result[0]
        for r_ind, r in enumerate(res):
            det_cnt += len(r)
            for rr_ind, rr in enumerate(r):
                if 'det_type' in rr and rr['det_type'] == 'det_fn':
                    print("DETECTED FN IN test_results_list")
                    print(rr['det_type'])                
                if not 'box_score' in rr.keys():
                    continue
                else:
                    pid_ixs = (test_df['pid'] == name)
                    pred_score_ixs = (test_df['pred_score'] == rr['box_score'])
                    det_type = test_df[pid_ixs & pred_score_ixs]['det_type'].tolist()[0]
                    test_results_list[test_result_ind][0][r_ind][rr_ind].update({'det_type' : det_type})

    return test_results_list

def save_charactersitic_plots():

    exp_dict = {
        'high_nonlesion' : 'high confidence non lesions (false positives)',
        'low_lesion' : 'low confidence lesions (true positives)'
    }
    metric_dict = {
        'dims' : 'Dimensions',
        'centroid' : 'Positions'
    }

    df_path = os.path.join(cf.exp_dir, 'les.pkl')
    df = pd.read_pickle(df_path)
    plot_dir = os.path.join(cf.exp_dir, 'plots', 'char_plots')
    if not os.path.exists(plot_dir):
        os.mkdir(plot_dir)

    for comment in np.unique(df['comment']):
        if 'good' in comment:
            continue
        wbc_str = 'without_wbc' if 'without' in comment else 'with_wbc'
        exp_string = '_'.join(comment.split('_')[:2])
        dff = df[df['comment'] == comment]
        dffg = df[df['comment'] == 'good_tp_' + wbc_str]       
        for metric in ['dims', 'centroid']:  
            print(metric, comment, wbc_str) 
            f, axs = plt.subplots(1, 3, figsize = (20, 5))
            fig_title = metric_dict[metric] + ' of ' + exp_dict[exp_string] + ' ' + ' '.join(wbc_str.split('_'))
            for ind, (dim, ax) in enumerate(zip(['width', 'height', 'depth'], axs)):
                d = dff[metric].apply(lambda x : x[ind])
                dg = dffg[metric].apply(lambda x : x[ind])
                plt.subplot(1, 3, ind+1)
                plt.hist([d, dg], bins = 10, label=[exp_dict[exp_string], 'high confidence lesions (true positives)'])
                # ax.set_title(dim)
                ax.set_xlabel(dim)
                ax.set_ylabel('Frequency')
                plt.legend(loc='upper right')
            f.suptitle(fig_title)
            plt.savefig(os.path.join(plot_dir, fig_title + '.png'), facecolor="white")   
            # # plt.show()
            # break

def get_lesions():
    wbc_str = 'with_wbc'
    fold_df_paths = [ii for ii in os.listdir(cf.exp_dir) if 'test_df_{}.pickle'.format(wbc_str) in ii]
    print(fold_df_paths)
    fold = cf.fold
    if len(fold_df_paths) == cf.n_cv_splits:
        dfs_list = [pd.read_pickle(os.path.join(cf.exp_dir, ii)) for ii in fold_df_paths]
        for ix, df in enumerate(dfs_list):
            df['fold'] = ix
        test_df = pd.concat(dfs_list)
    print(test_df.head())
    print(test_df.shape)
    print(np.unique(test_df['fold'], return_counts=True))
    print(test_df['pid'].value_counts())

if __name__ == '__main__':

    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_to_checkpoint', type=str, default=None,
                        help='if resuming to checkpoint, the desired fold still needs to be parsed via --folds.')
    parser.add_argument('--exp_source', type=str, default='/path/to/experiment/source',
                        help='specifies, from which source experiment to load configs and data_loader.')
    parser.add_argument('--pre_train_path',type=str, default=None, 
                        help="full path to pretrained checkpoint")
    parser.add_argument('-d', '--dev', default=False, action='store_true', help="development mode: shorten everything")


    args = parser.parse_args()
    folds = args.folds
    resume_to_checkpoint = args.resume_to_checkpoint
    pre_train_path = args.pre_train_path
    
    # setting device on GPU if available, else CPU
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type == 'cuda' : print(torch.cuda.get_device_name(0))
    
    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]
            cf.n_workers = 1
            cf.select_prototype_subset = 6
            cf.batch_size, cf.num_epochs, cf.min_save_thresh, cf.save_n_models = 3 if cf.dim==2 else 2, 2, 0, 2
            cf.num_train_batches, cf.num_val_batches, cf.max_val_patients = 2, 2, 2
            cf.test_n_epochs =  cf.save_n_models
            cf.max_test_patients = 1

        if args.data_dest:
            cf.data_dest = args.data_dest
        
        model = utils.import_module('model', cf.model_path)
        data_loader = utils.import_module('dl', os.path.join(args.exp_dir, 'data_loader.py'))
        cf.folds = range(cf.n_cv_splits)
        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
            cf.resume_to_checkpoint = resume_to_checkpoint
            cf.pre_train_path = pre_train_path
            if not os.path.exists(cf.fold_dir):
                os.mkdir(cf.fold_dir)
            logger = utils.get_logger(cf.fold_dir)
            train(logger)
            cf.resume_to_checkpoint = None
            if args.mode == 'train_test':
                test(logger)

            for hdlr in logger.handlers:
                hdlr.close()
            logger.handlers = []

    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 =  1; cf.max_test_patients = 1

        if args.data_dest:
            cf.data_dest = args.data_dest
        model = utils.import_module('model', cf.model_path)
        data_loader = utils.import_module('dl', os.path.join(args.exp_dir, 'data_loader.py'))
        if folds is None:
            folds = range(cf.n_cv_splits)
        cf.folds = folds
        for fold in folds:
            cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))
            logger = utils.get_logger(cf.fold_dir)
            cf.fold = fold
            print("Fold" , cf.fold)
            test(logger)

            for hdlr in logger.handlers:
                hdlr.close()
            logger.handlers = []
    
    elif args.mode == 'lcl':
        cf = utils.prep_exp(args.exp_source, args.exp_dir, args.server_env, is_training=False, use_stored_settings=True)
        if args.data_dest:
            cf.data_dest = args.data_dest
        model = utils.import_module('model', cf.model_path)
        data_loader = utils.import_module('dl', os.path.join(args.exp_dir, 'data_loader.py'))
        if folds is None:
            folds = range(cf.n_cv_splits)
        les_df_path = os.path.join(cf.exp_dir, 'les.pkl')
        bool_clear = False
        if os.path.exists(les_df_path):
            if bool_clear:
                os.remove(les_df_path)
            else:
                all_df = pd.read_pickle(les_df_path)
        if not os.path.exists(les_df_path):
            cf.folds = folds
            cols = ['p_id', 'box_type', 'box_coords', 'box_score', 'box_pred_class_id', 'det_type', 'dims', 'centroid', 'comment', 'fold']
            all_df = pd.DataFrame(columns=cols)      
            for fold in folds:
                cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))
                logger = utils.get_logger(cf.fold_dir)
                cf.fold = fold
                print("Fold" , cf.fold)
                df = lcl(logger)
                df['fold'] = cf.fold
                all_df = all_df.append(df, ignore_index=True)
                print(all_df.shape)
                for hdlr in logger.handlers:
                    hdlr.close()
                logger.handlers = []
            all_df.to_pickle(les_df_path)
        print(all_df['comment'].value_counts())
        save_charactersitic_plots()

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

        if cf.hold_out_test_set:
            data_loader = utils.import_module('dl', os.path.join(args.exp_dir, 'data_loader.py'))
            cf.folds = args.folds
            predictor = Predictor(cf, net=None, logger=logger, mode='analysis')
            results_list = predictor.load_saved_predictions(apply_wbc=True)
            utils.create_csv_output(results_list, cf, logger)

        else:
            if folds is None:
                folds = range(cf.n_cv_splits)
            for fold in folds:
                print("="*10, fold)
                cf.fold_dir = os.path.join(cf.exp_dir, 'fold_{}'.format(fold))
                cf.fold = fold
                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)
                print(len(evaluator.test_df['det_type']))
                evaluator.score_test_df()

    # 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=True)
        logger = utils.get_logger(cf.exp_dir)
        logger.info('created experiment directory at {}'.format(args.exp_dir))

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