import pandas as pd

import numpy as np

import matplotlib.pyplot as plt

plt.switch_backend('agg')

import seaborn as sns

sns.set_style("whitegrid")

import math

import argparse

from glob import glob

import pdb

def main():

    """Analyzes results and generates figures."""

    parser = argparse.ArgumentParser(description="An analyst for quick ML applications.",

                                     add_help=False)

    parser.add_argument('RUN_DIR', action='store',  type=str, help='Path to results from analysis.')    

    parser.add_argument('-max_feat',action='store',dest='MAX_FEAT',default=10,type=int,

                        help = 'Max features to show in importance plots.')

    args = parser.parse_args()

    # dataset = args.NAME

    # dataset = args.NAME.split('/')[-1].split('.')[0] 

    # run_dir = 'results/' + dataset + '/' 

    run_dir = args.RUN_DIR

    if run_dir[-1] != '/': 

        run_dir += '/'

    dataset = run_dir.split('/')[-2]

    print('dataset:',dataset)

    print('loading data from',run_dir)

    frames = []     # data frames to combine

    count = 0

    for f in glob(run_dir + '*.csv'):

        if 'imp_score' not in f:

            frames.append(pd.read_csv(f,sep='\t',index_col=False))

            count = count + 1

    df = pd.concat(frames, join='outer', ignore_index=True)

    print('loaded',count,'result files with results from these learners:',df['algorithm'].unique())

    restricted_cols = ['prep_alg','preprocessor', 'prep-parameters', 'algorithm', 'alg-parameters','dataset',

                      'trial','seed','parameters']

    columns_to_plot = [c for c in df.columns if c not in restricted_cols ] 

    #['accuracy','f1_macro','bal_accuracy']

    print('generating boxplots for these columns:',columns_to_plot)

    for col in columns_to_plot:

        fig = plt.figure()

        # for i, prep in enumerate(unique_preps):

            # fig.add_subplot(math.ceil(len(unique_preps)), 2,i+1)

        # pdb.set_trace()

        df[col] = df[col].astype(np.float)

        sns.boxplot(data=df,x="algorithm",y=col)

        # plt.title(prep,size=16)

        plt.gca().set_xticklabels(df.algorithm.unique(),size=14,rotation=45)

        plt.ylabel(col,size=16)

        plt.ylim(0.5,1.0)

        plt.xlabel('')

        fig.tight_layout() 

        plt.savefig(run_dir + '_'.join([ dataset, col,'boxplots.pdf']))

    ####################################################################### feature importance plots

    frames = []     # data frames to combine

    count = 0

    for f in glob(run_dir + '*.imp_score'):

        frames.append(pd.read_csv(f,sep='\t',index_col=False))

        count = count + 1

    df = pd.concat(frames, join='outer', ignore_index=True)

    print('loaded',count,'feature importance files with results from these learners:',df['algorithm'].unique())

    dfp =  df.groupby(['algorithm','feature']).median().unstack(['algorithm'])

    dfpn = df.groupby(['feature','algorithm']).median().groupby('feature').sum().unstack()

    dfpn.sort_values(ascending=False, inplace=True)

    # sort by median feature importance

    nf = min(args.MAX_FEAT, dfpn.index.labels[1].shape[0])

    dfpw = dfp.iloc[dfpn.index.labels[1][:nf]]

    h = dfpw['score'].plot(kind='bar', stacked=True)

    leg = plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))

    plt.ylabel('Importance Score')

    plt.savefig(run_dir + '_'.join([ dataset, 'importance_scores.pdf']),bbox_extra_artists=(leg,h), bbox_inches='tight')

    ############################################################# roc curves

    frames = []     # data frames to combine

    count = 0

    for f in glob(run_dir + '*.roc'):

        frames.append(pd.read_csv(f,sep='\t',index_col=False))

        count = count + 1

    df = pd.concat(frames, join='outer', ignore_index=True)

    print('loaded',count,'roc files with results from these learners:',df['algorithm'].unique())

    h, ax = plt.subplots()

    ax.plot([0, 1],[0, 1],'--k',label='_nolegend_')

    colors = ('r','y','b','g','c','k')

    colors = plt.cm.Blues(np.linspace(0.1, 0.9, len(df['algorithm'].unique())))

    n_algs = len(df['algorithm'].unique())

    markers = ['o','v','^','<','>','8','s',

               'p','P','*','h','H','+','x','X','D','d','|','_']

    for i, (alg,df_g) in enumerate(df.groupby('algorithm')):

        aucs = df_g.auc.values

        seed_max = df_g.loc[df_g.auc.idxmax()]['seed']

        seed_min = df_g.loc[df_g.auc.idxmin()]['seed']

        seed_med = df_g.loc[np.abs(df_g.auc - df_g.auc.median()) == np.min(np.abs(df_g.auc - df_g.auc.median()))]['seed']

        seed_med = seed_med.iloc[0]

        auc = df_g.auc.median()

        # fpr = df_g['fpr'].unique()

        tprs,fprs=[],[]

        fpr_min = df_g.loc[df_g.seed == seed_min,:]['fpr']

        fpr_max = df_g.loc[df_g.seed == seed_max,:]['fpr']

        tpr_min = df_g.loc[df_g.seed == seed_min,:]['tpr']

        tpr_max = df_g.loc[df_g.seed == seed_max,:]['tpr']

        tpr_med = df_g.loc[df_g.seed == seed_med,:]['tpr']

        fpr_med = df_g.loc[df_g.seed == seed_med,:]['fpr']

        ax.plot(fpr_med,tpr_med, color=colors[i % n_algs], marker=markers[i], 

                linestyle='--', linewidth=1, label='{:s} (AUC = {:0.2f})'.format(alg,auc))

        # ax.plot(fpr_max,tpr_max, color=colors[i % n_algs],  

        #         linestyle='--', linewidth=1, label='_nolegend_', alpha=0.1)

        # ax.plot(fpr_min,tpr_min,color=colors[i % n_algs],  

        #         linestyle='--', linewidth=1, label='_nolegend_', alpha=0.1)

    plt.ylabel('True Positive Rate')

    plt.xlabel('False Positive Rate')

    leg = plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))

    plt.ylim(0,1)

    plt.xlim(0,1)

    plt.savefig(run_dir + '_'.join([ dataset, 'roc_curves.pdf']), bbox_inches='tight')

    print('done!')    

if __name__ == '__main__':

    main()