#evaluate.py

#Copyright (c) 2020 Rachel Lea Ballantyne Draelos

#MIT License

#Permission is hereby granted, free of charge, to any person obtaining a copy

#of this software and associated documentation files (the "Software"), to deal

#in the Software without restriction, including without limitation the rights

#to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

#copies of the Software, and to permit persons to whom the Software is

#furnished to do so, subject to the following conditions:

#The above copyright notice and this permission notice shall be included in all

#copies or substantial portions of the Software.

#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

#FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

#AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

#LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

#OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE

#SOFTWARE

#Imports

import os

import copy

import time

import bisect

import shutil

import operator

import itertools

import numpy as np

import pandas as pd

import sklearn.metrics

from scipy import interp

from itertools import cycle

import matplotlib

matplotlib.use('agg') #so that it does not attempt to display via SSH

import seaborn

import matplotlib.pyplot as plt

plt.ioff() #turn interactive plotting off

#suppress numpy warnings

import warnings

warnings.filterwarnings('ignore')

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

# Reporting Functions #---------------------------------------------------------

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

def initialize_evaluation_dfs(all_labels, num_epochs):

    """Create empty "eval_dfs_dict"

    Variables

    <all_labels>: a list of strings describing the labels in order

    <num_epochs>: int for total number of epochs"""

    if len(all_labels)==2:

        index = [all_labels[1]]

        numrows = 1

    else:

        index = all_labels

        numrows = len(all_labels)

    #Initialize empty pandas dataframe to store evaluation results across epochs

    #for accuracy, AUROC, and AP

    result_df = pd.DataFrame(data=np.zeros((numrows, num_epochs)),

                            index = index,

                            columns = ['epoch_'+str(n) for n in range(0,num_epochs)])

    #Initialize empty pandas dataframe to store evaluation results for top k

    top_k_result_df = pd.DataFrame(np.zeros((len(all_labels), num_epochs)),

                                   index=[x for x in range(1,len(all_labels)+1)], #e.g. 1,...,64 for len(all_labels)=64

                                   columns = ['epoch_'+str(n) for n in range(0,num_epochs)])

    #Make eval results dictionaries

    eval_results_valid = {'accuracy':copy.deepcopy(result_df),

        'auroc':copy.deepcopy(result_df),

        'avg_precision':copy.deepcopy(result_df),

        'top_k':top_k_result_df}

    eval_results_test = copy.deepcopy(eval_results_valid)

    return eval_results_valid, eval_results_test

def save(eval_dfs_dict, results_dir, descriptor):

    """Variables

    <eval_dfs_dict> is a dict of pandas dataframes

    <descriptor> is a string"""

    for k in eval_dfs_dict.keys():

        eval_dfs_dict[k].to_csv(os.path.join(results_dir, descriptor+'_'+k+'_Table.csv'))

def save_final_summary(eval_dfs_dict, best_valid_epoch, setname, results_dir):

    """Save to overall df and print summary of best epoch."""

    #final_descriptor is e.g. '2019-11-15-awesome-model_epoch15

    final_descriptor = results_dir.replace('results/','')+'_epoch'+str(best_valid_epoch)

    if setname=='valid': print('***Summary for',setname,results_dir,'***')

    for metricname in list(eval_dfs_dict.keys()):

        #metricnames are accuracy, auroc, avg_precision, and top_k.

        #df holds a particular metric for the particular model we just ran.

        #for accuracy, auroc, and avg_precision, df index is diseases, columns are epochs.

        #for top_k, df index is the k value (an int) and columns are epochs.

        df = eval_dfs_dict[metricname]

        #all_df tracks results of all models in one giant table.

        #all_df has index of diseases or k value, and columns which are particular models.

        all_df_path = os.path.join('results',setname+'_'+metricname+'_all.csv') #e.g. valid_accuracy_all.csv

        if os.path.isfile(all_df_path):

            all_df = pd.read_csv(all_df_path,header=0,index_col=0)

            all_df[final_descriptor] = np.nan

        else: #all_df doesn't exist yet - create it.

            all_df = pd.DataFrame(np.empty((df.shape[0],1)),

                                  index = df.index.values.tolist(),

                                  columns = [final_descriptor])

        #Print off and save results for best_valid_epoch

        if setname=='valid': print('\tEpoch',best_valid_epoch,metricname)

        for label in df.index.values:

            #print off to console

            value = df.at[label,'epoch_'+str(best_valid_epoch)]

            if setname=='valid': print('\t\t',label,':',str( round(value, 3) ))

            #save in all_df

            all_df.at[label,final_descriptor] = value

        all_df.to_csv(all_df_path,header=True,index=True)

def clean_up_output_files(best_valid_epoch, results_dir):

    """Delete output files that aren't from the best epoch"""

    #Delete all the backup parameters (they take a lot of space and you do not

    #need to have them)

    shutil.rmtree(os.path.join(results_dir,'backup'))

    #Delete all the extra output files:

    for subdir in ['heatmaps','curves','pred_probs']:

        #Clean up saved ROC and PR curves

        fullpath = os.path.join(results_dir,subdir)

        if os.path.exists(fullpath): #e.g. there may not be a heatmaps dir for a non-bottleneck model

            allfiles = os.listdir(fullpath)

            for filename in allfiles:

                if str(best_valid_epoch) not in filename:

                    os.remove(os.path.join(fullpath,filename))

    print('Output files all clean')

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

# Calculation Functions #-------------------------------------------------------

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

def evaluate_all(eval_dfs_dict, epoch, label_meanings,

                 true_labels_array, pred_probs_array):

    """Fill out the pandas dataframes in the dictionary <eval_dfs_dict>

    which is created in cnn.py. <epoch> and <which_label> are used to index into

    the dataframe for the metric. Metrics calculated for the provided vectors

    are: accuracy, AUC, partial AUC (threshold 0.2), and average precision.

    If <subjective> is set to True, additional metrics will be calculated

    (confusion matrix, sensitivity, specificity, PPV, NPV.)

    Variables:

    <all_eval_results> is a dictionary of pandas dataframes created in cnn.py

    <epoch> is an integer indicating which epoch it is, starting from epoch 1

    <true_labels_array>: array of true labels. examples x labels

    <pred_probs_array>: array of predicted probabilities. examples x labels"""

    #Accuracy, AUROC, and AP (iter over labels)

    for label_number in range(len(label_meanings)):

        which_label = label_meanings[label_number] #descriptive string for the label

        true_labels = true_labels_array[:,label_number]

        pred_probs = pred_probs_array[:,label_number]

        pred_labels = (pred_probs>=0.5).astype(dtype='int') #decision threshold of 0.5

        #Accuracy and confusion matrix (dependent on decision threshold)

        (eval_dfs_dict['accuracy']).at[which_label, 'epoch_'+str(epoch)] = compute_accuracy(true_labels, pred_labels)

        #confusion_matrix, sensitivity, specificity, ppv, npv = compute_confusion_matrix(true_labels, pred_labels)

        #AUROC and AP (sliding across multiple decision thresholds)

        fpr, tpr, thresholds = sklearn.metrics.roc_curve(y_true = true_labels,

                                         y_score = pred_probs,

                                         pos_label = 1)

        (eval_dfs_dict['auroc']).at[which_label, 'epoch_'+str(epoch)] = sklearn.metrics.auc(fpr, tpr)

        (eval_dfs_dict['avg_precision']).at[which_label, 'epoch_'+str(epoch)] = sklearn.metrics.average_precision_score(true_labels, pred_probs)

    #Top k eval metrics (iter over examples)

    eval_dfs_dict['top_k'] = evaluate_top_k(eval_dfs_dict['top_k'],

                                    epoch, true_labels_array, pred_probs_array)

    return eval_dfs_dict

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

# Top K Metrics #---------------------------------------------------------------

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

def evaluate_top_k(eval_top_k_df, epoch, true_labels_array,

                   pred_probs_array):

    """<eval_top_k_df> is a pandas dataframe with epoch number as columns and

        k values as rows, where k is an integer"""

    num_labels = true_labels_array.shape[1] #e.g. 64

    total_examples = true_labels_array.shape[0]

    vals = [0 for x in range(1,num_labels+2)] #e.g. length 65 list but the index of the last element is 64 for num_labels=64

    for example_number in range(total_examples):

        #iterate through individual examples (predictions for an individual CT)

        #rather than iterating through predicted labels

        true_labels = true_labels_array[example_number,:]

        pred_probs = pred_probs_array[example_number,:]

        for k in range(1,num_labels+1): #e.g. 1,...,64

            previous_value = vals[k]

            incremental_update = calculate_top_k_accuracy(true_labels, pred_probs, k)

            new_value = previous_value + incremental_update

            vals[k] = new_value

    #Now update the dataframe. Should reach 100% performance by the end.

    for k in range(1,num_labels+1):

        eval_top_k_df.at[k,'epoch_'+str(epoch)] = vals[k]/total_examples

    ##Now average over all the examples

    #eval_top_k_df.loc[:,'epoch_'+str(epoch)] = eval_top_k_df.loc[:,'epoch_'+str(epoch)] / total_examples

    return eval_top_k_df

def calculate_top_k_accuracy(true_labels, pred_probs, k):

    k = min(k, len(true_labels)) #avoid accessing array elements that don't exist

    #argpartition described here: https://stackoverflow.com/questions/6910641/how-do-i-get-indices-of-n-maximum-values-in-a-numpy-array

    #get the indices of the largest k probabilities

    ind = np.argpartition(pred_probs, -1*k)[-1*k:]

    #now figure out what percent of these top predictions were equal to 1 in the

    #true_labels.

    #Note that the denominator should not exceed the number of true labels, to

    #avoid penalizing the model inappropriately:

    denom = min(k, np.sum(true_labels))

    if denom == 0: #because np.sum(true_labels) is 0

        #super important! must return 1 to avoid dividing by 0 and producing nan

        #we don't return 0 because then the model can never get perfect performance

        #even at k=num_labels because it'll get 0 for anything that has no labels

        return 1 

    else:

        return float(np.sum(true_labels[ind]))/denom

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

# Accuracy and AUROC #----------------------------------------------------------

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

def compute_accuracy(true_labels, labels_pred):

    """Print and save the accuracy of the model on the dataset"""    

    correct = (true_labels == labels_pred)

    correct_sum = correct.sum()

    return (float(correct_sum)/len(true_labels))

def compute_confusion_matrix(true_labels, labels_pred):

    """Return the confusion matrix"""

    cm = sklearn.metrics.confusion_matrix(y_true=true_labels,

                          y_pred=labels_pred)

    if cm.size < 4: #cm is too small to calculate anything

        return np.nan, np.nan, np.nan, np.nan, np.nan

    true_neg, false_pos, false_neg, true_pos = cm.ravel()

    sensitivity = float(true_pos)/(true_pos + false_neg)

    specificity = float(true_neg)/(true_neg + false_pos)

    ppv = float(true_pos)/(true_pos + false_pos)

    npv = float(true_neg)/(true_neg + false_neg)

    return((str(cm).replace("\n","_")), sensitivity, specificity, ppv, npv)

def compute_partial_auroc(fpr, tpr, thresh = 0.2, trapezoid = False, verbose=False):

    fpr_thresh, tpr_thresh = get_fpr_tpr_for_thresh(fpr, tpr, thresh)

    if len(fpr_thresh) < 2:#can't calculate an AUC with only 1 data point

        return np.nan 

    if verbose:

        print('fpr: '+str(fpr))

        print('fpr_thresh: '+str(fpr_thresh))

        print('tpr: '+str(tpr))

        print('tpr_thresh: '+str(tpr_thresh))

    return sklearn.metrics.auc(fpr_thresh, tpr_thresh)

def get_fpr_tpr_for_thresh(fpr, tpr, thresh):

    """The <fpr> and <tpr> are already sorted according to threshold (which is

    sorted from highest to lowest, and is NOT the same as <thresh>; threshold

    is the third output of sklearn.metrics.roc_curve and is a vector of the

    thresholds used to calculate FPR and TPR). This function figures out where

    to bisect the FPR so that the remaining elements are no greater than

    <thresh>. It bisects the TPR in the same place."""

    p = (bisect.bisect_left(fpr, thresh)-1) #subtract one so that the FPR

    #of the remaining elements is NO GREATER THAN <thresh>

    return fpr[: p + 1], tpr[: p + 1]

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

# Plotting Functions #----------------------------------------------------------

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

def plot_pr_and_roc_curves(results_dir, label_meanings, true_labels, pred_probs,

                           epoch):

    #Plot Precision Recall Curve

    #Plot ROC Curve

    fpr, tpr, thresholds = sklearn.metrics.roc_curve(y_true = true_labels,

                                         y_score = pred_probs,

                                         pos_label = 1)

    plot_roc_curve(fpr, tpr, epoch, outfilepath)

def plot_roc_curve_multi_class(label_meanings, y_test, y_score, 

                               outdir, setname, epoch):

    """<label_meanings>: list of strings, one for each label

    <y_test>: matrix of ground truth

    <y_score>: matrix of predicted probabilities

    <outdir>: directory to save output file

    <setname>: string e.g. 'train' 'valid' or 'test'

    <epoch>: int for epoch"""

    #Modified from https://scikit-learn.org/stable/auto_examples/model_selection/plot_roc.html

    n_classes = len(label_meanings)

    lw = 2

    # Compute ROC curve and ROC area for each class

    fpr = dict()

    tpr = dict()

    roc_auc = dict()

    for i in range(n_classes):

        fpr[i], tpr[i], _ = sklearn.metrics.roc_curve(y_test[:, i], y_score[:, i])

        roc_auc[i] = sklearn.metrics.auc(fpr[i], tpr[i])

    #make order df. (note that roc_auc is a dictionary with ints as keys

    #and AUCs as values. 

    order = pd.DataFrame(np.zeros((n_classes,1)), index = [x for x in range(n_classes)],

                         columns = ['roc_auc'])

    for i in range(n_classes):

        order.at[i,'roc_auc'] = roc_auc[i]

    order = order.sort_values(by='roc_auc',ascending=False)

    #Plot all ROC curves

    #Plot in order of the rainbow colors, from highest AUC to lowest AUC

    plt.figure()

    colors_list = ['palevioletred','darkorange','yellowgreen','olive','deepskyblue','royalblue','navy']

    curves_plotted = 0

    for i in order.index.values.tolist()[0:10]: #only plot the top ten so the plot is readable

        color_idx = curves_plotted%len(colors_list) #cycle through the colors list in order of colors

        color = colors_list[color_idx]

        plt.plot(fpr[i], tpr[i], color=color, lw=lw,

                 label='{:5s} (area {:0.2f})'.format(label_meanings[i], roc_auc[i]))

        curves_plotted+=1

    plt.plot([0, 1], [0, 1], 'k--', lw=lw)

    plt.xlim([0.0, 1.0])

    plt.ylim([0.0, 1.05])

    plt.xlabel('False Positive Rate')

    plt.ylabel('True Positive Rate')

    plt.title(setname.lower().capitalize()+' ROC Epoch '+str(epoch))

    plt.legend(loc="lower right",prop={'size':6})

    outfilepath = os.path.join(outdir,setname+'_ROC_ep'+str(epoch)+'.pdf')

    plt.savefig(outfilepath)

    plt.close()

def plot_pr_curve_multi_class(label_meanings, y_test, y_score, 

                              outdir, setname, epoch):

    """<label_meanings>: list of strings, one for each label

    <y_test>: matrix of ground truth

    <y_score>: matrix of predicted probabilities

    <outdir>: directory to save output file

    <setname>: string e.g. 'train' 'valid' or 'test'

    <epoch>: int for epoch"""

    #Modified from https://scikit-learn.org/stable/auto_examples/model_selection/plot_roc.html

    #https://stackoverflow.com/questions/29656550/how-to-plot-pr-curve-over-10-folds-of-cross-validation-in-scikit-learn

    n_classes = len(label_meanings)

    lw = 2

    #make order df.

    order = pd.DataFrame(np.zeros((n_classes,1)), index = [x for x in range(n_classes)],

                         columns = ['prc'])

    for i in range(n_classes):

        order.at[i,'prc'] = sklearn.metrics.average_precision_score(y_test[:,i], y_score[:,i])

    order = order.sort_values(by='prc',ascending=False)

    #Plot

    plt.figure()

    colors_list = ['palevioletred','darkorange','yellowgreen','olive','deepskyblue','royalblue','navy']

    curves_plotted = 0

    for i in order.index.values.tolist()[0:10]: #only plot the top ten so the plot is readable

        color_idx = curves_plotted%len(colors_list) #cycle through the colors list in order of colors

        color = colors_list[color_idx]

        average_precision = sklearn.metrics.average_precision_score(y_test[:,i], y_score[:,i])

        precision, recall, _ = sklearn.metrics.precision_recall_curve(y_test[:,i], y_score[:,i])

        plt.step(recall, precision, color=color, where='post',

                 label='{:5s} (area {:0.2f})'.format(label_meanings[i], average_precision))

        curves_plotted+=1

    plt.xlabel('Recall')

    plt.ylabel('Precision')

    plt.ylim([0.0, 1.05])

    plt.xlim([0.0, 1.0])

    plt.title(setname.lower().capitalize()+' PRC Epoch '+str(epoch))

    plt.legend(loc="lower right",prop={'size':6})

    outfilepath = os.path.join(outdir,setname+'_PR_ep'+str(epoch)+'.pdf')

    plt.savefig(outfilepath)

    plt.close()

def plot_pr_curve_single_class(true_labels, pred_probs, outfilepath):

    #http://scikit-learn.org/stable/auto_examples/model_selection/plot_precision_recall.html

    average_precision = sklearn.metrics.average_precision_score(true_labels, pred_probs)

    precision, recall, _ = sklearn.metrics.precision_recall_curve(true_labels, pred_probs)

    plt.step(recall, precision, color='b', alpha=0.2,

             where='post')

    plt.fill_between(recall, precision, step='post', alpha=0.2,

                     color='b')

    plt.xlabel('Recall')

    plt.ylabel('Precision')

    plt.ylim([0.0, 1.05])

    plt.xlim([0.0, 1.0])

    plt.title('2-class Precision-Recall curve: AP={0:0.2f}'.format(

              average_precision))

    plt.savefig(outfilepath)

    plt.close()

def plot_roc_curve_single_class(fpr, tpr, epoch, outfilepath):

    #http://scikit-learn.org/stable/auto_examples/model_selection/plot_roc.html#sphx-glr-auto-examples-model-selection-plot-roc-py

    roc_auc = sklearn.metrics.auc(fpr, tpr)

    lw = 2

    plt.plot(fpr, tpr, color='darkorange',

             lw=lw, label='ROC curve (area = %0.2f)' % roc_auc)

    plt.plot([0, 1], [0, 1], color='navy', lw=lw, linestyle='--')

    plt.xlim([0.0, 1.0])

    plt.ylim([0.0, 1.05])

    plt.xlabel('False Positive Rate')

    plt.ylabel('True Positive Rate')

    plt.title('Receiver Operating Characteristic')

    plt.legend(loc="lower right")

    plt.savefig(outputfilepath)

    plt.close()

def plot_learning_curves(train_loss, valid_loss, results_dir, descriptor):

    """Variables

    <train_loss> and <valid_loss> are numpy arrays with one numerical entry

    for each epoch quanitfying the loss for that epoch."""

    x = np.arange(0,len(train_loss))

    plt.plot(x, train_loss, color='blue', lw=2, label='train')

    plt.plot(x, valid_loss, color='green',lw = 2, label='valid')

    plt.xlabel('Epoch')

    plt.ylabel('Loss')

    plt.title('Training and Validation Loss')

    plt.legend(loc='lower right')

    plt.savefig(os.path.join(results_dir, descriptor+'_Learning_Curves.png'))

    plt.close()

    #save numpy arrays of the losses

    np.save(os.path.join(results_dir,'train_loss.npy'),train_loss)

    np.save(os.path.join(results_dir,'valid_loss.npy'),valid_loss)

def plot_heatmap(outprefix, numeric_array, center, xticklabels, yticklabels):

    """Save a heatmap based on numeric_array"""

    seaborn.set(font_scale=0.6)

    seaplt = (seaborn.heatmap(numeric_array,

                           center=center,

                           xticklabels=xticklabels,

                           yticklabels=yticklabels)).get_figure()

    seaplt.savefig(outprefix+'.png')

    seaplt.clf()