stats.py

## Essentially intended to be a way to create some nice stats plots for the paper
## July 2021-July 2022: 1807 candidates, average of 4.36 candidates per day

## Need to stylize the plots; could define a function to outline a consistent plot style and size depending on plot type to reduce repetition - see https://stackoverflow.com/questions/51711438/matplotlib-how-to-edit-the-same-plot-with-different-functions and https://matplotlib.org/1.5.3/users/style_sheets.html for more info
## need to fig,ax stuff

## could probably consolidate some of the plotting functions into one function for each topic (e.g. one function for plotting the number of fits, one for plotting the number of unfit, etc.)

## alternatively, could change these into methods of a class that takes the dataframe as the object (e.g. df.plotDailyCand() ) 

import argparse
import glob
import itertools
import json
import os
import sys
import time
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)


import matplotlib as mpl
import matplotlib.dates as dates
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import pandas as pd
import seaborn as sns

from astropy.time import Time

## to prevent a bunch of messages when defining df_fo
pd.options.mode.chained_assignment = None

#from scipy.interpolate import make_interp_spline as spline

## set plot style
plt.style.use("seaborn-colorblind")
mpl.rcParams.update({"axes.grid" : True})
plt.style.context(("seaborn-colorblind",))
# mpl.rcParams['text.usetex'] = True
# mpl.rcParams['text.latex.preamble'] = [r'\usepackage{amsmath}']

## argument for folder to run stats on
parser = argparse.ArgumentParser()
parser.add_argument("-c","--candDir", type=str, default=None, help="Path to the candidate directory")
parser.add_argument("-f","--fitDir", type=str, default=None, help="Path to the fits directory")
parser.add_argument('-v', '--verbose', action='store_true')
parser.add_argument('-o', '--outdir', type=str, default=os.path.join('./outdir/stats/',time.strftime("%Y%m%d-%H%M%S"),''))
parser.add_argument("-m","--models", nargs="+", default = ["nugent-hyper","Bu2019lm","TrPi2018","Piro2021"], choices = ["TrPi2018","nugent-hyper", "Piro2021","Bu2019lm"], help="which models to analyse with the fit stats")
parser.add_argument('-df','--datafile',type=str, default=None, help="path to the csv file that's generated by the stats.py script")
args = parser.parse_args()

## to change to correct directory
os.chdir(sys.path[0])
print("Current working directory: {}\n".format(os.getcwd())) if args.verbose else None

## compilation of lists for use in plotting (pre-dataframe implementation)
## post dataframe implementation: these should eventually be removed and plots should be updated to use the dataframe
if args.candDir:
    dayList = glob.glob(os.path.join(args.candDir, "/*/"))
    
    dayCount = [day for day in range(0,len(dayList))]
    
    numDaily = [len(glob.glob(day + "/*.csv")) for day in dayList]
    
    candList = glob.glob(args.candDir + "*/*.csv")

    cumDaily = np.cumsum(numDaily)
else:
    print("No candidate directory specified, cannot run some stats\n")

if args.fitDir:
    fit_dayList = glob.glob(args.fitDir + "*")
    
    ## attempted dict comprehension for finding all instances of logs
    logDict = {model: glob.glob(os.path.join(args.fitDir,'*',model+'.log')) for model in args.models}
    jsonDict = {model: glob.glob(os.path.join(args.fitDir,'*',model+'_result.json')) for model in args.models}
    ## the way these are written, it may have issues for plotting if some days don't have all models
else:
    print("No candidate directory specified, cannot run some stats\n")


## Utility functions
def plotDir(name,outdir=args.outdir,ext=".png"): ## might be good to organize different plot types into subdirectories, but doesn't have to be an argument here
    '''
    check for existence of plot directory and create if needed, then return full path for saving figure
    
    Args:
    name: name of the plot for the filename (without extension)
    outdir: path to the output directory
    ext: extension of the plot file (typically .png or .pdf)
    '''
    if not os.path.exists(outdir):
        os.makedirs(outdir)
    filepath = os.path.join(outdir,name+ext)
    return(filepath)


## function is currently unused/incomplete, wanted to at least start it and make a commit with it
def plotstyle(type=None, **kwargs): ## should add an option to pass kwargs to the plot function
    '''
    Sets the style of the plots to be consistent across the paper using matplotlib's style sheets
    
    Args:
    type: type of plot to be generated (e.g. 'histogram', 'scatter')
    **kwargs: anything you would pass to matplotlib
    '''
    fig, ax = plt.subplots(**kwargs) ## not super sure on this implementation
    
    plt.style.use('seaborn-whitegrid')
    plt.style.context(("seaborn-colorblind",))
    #plt.rcParams['font.family'] = 'serif'
    #plt.rcParams['font.serif'] = 'Times New Roman'
    plt.rcParams['font.size'] = 24
    plt.rcParams['axes.labelsize'] = 24
    plt.rcParams['axes.labelweight'] = 'bold'
    plt.rcParams['axes.titlesize'] = 24
    plt.rcParams['xtick.labelsize'] = 24
    plt.rcParams['ytick.labelsize'] = 24
    plt.rcParams['legend.fontsize'] = 24
    plt.rcParams['figure.titlesize'] = 24

    return(fig,ax)


def get_sampling_time(file=None): ## somewhat redundant after creation of get_json
    '''
    Pulls from the provided json file to find the sampling_time and returns that value. sampling_time is recorded in seconds
    
    Args:
    file: path to the json file to be parsed
    '''
    if file:
        with open(file) as f:
            try: 
                data = json.load(f)
                sampling_time = data['sampling_time']
            finally:
                f.close()
                return sampling_time
            
    else:
        print('provide a file to search!\n')
        exit(1) ## irreconciable error, hence exit(1)


def get_json(file=None, params=None): ## effectively an improvement on get_sampling_time so it provides more flexibility and can be used for other parameters
    '''
    pulls from the provided json file to find several values and returns them in a dictionary. sampling_time is recorded in seconds. NOTE: all values are returned as strings and must be converted to the appropriate type when used.

    Args:
    file: path to json file (required). Taken as path string, but will also accept boolean False to return dictionary populated by np.nan values.
    params: list of additional parameters to pull from json file.
    '''
    jsonList = ['sampling_time', 'sampler', 'log_evidence','log_evidence_err', 'log_noise_evidence',  'log_bayes_factor']
    if params:
        jsonList = jsonList + params
    if file:
        try: 
            with open(file,'r') as f:
                data = json.load(f)
                jsonDict = {param: data[param] for param in jsonList}
                
        except: ## in event that json file isn't read correctly
            print('error reading json file: {} \n'.format(file))
            jsonDict = {param: np.nan for param in jsonList}
        finally:
            f.close()
            print('jsonDict: {}\n'.format(jsonDict)) if args.verbose else None
            return jsonDict
    elif not file: ## for use case where no json is found when the pandas dataframe is created
        jsonDict = {param: np.nan for param in jsonList} ## np.nan is used to make it easier to plot later without having to deal with NoneType
        print('no json file found/provided')
        print('jsonDict: {}\n'.format(jsonDict)) if args.verbose else None
        return jsonDict
    else: ## case where file argument is not provided
        print('provide a file to search!\n')
        exit(1) ## irreconciable error, hence exit(1)


def countDailyFits(day=None, models=args.models): ##relying on args as default might not be the best idea
    '''
    finds how many fits were completed on a given day, with day being provided as a path string
    Somewhat made redundant by the creation of get_dataframe, but I'll leave it for now
    Args:
    day: path to day directory
    models: list of models to search for
    '''
    if day:
        fitCands = glob.glob(os.path.join(day,'*/')) ## will return the paths to the candidates that were fit + the candidate_data folder 
        if os.path.join(day,'candidate_data/') in fitCands:
            candList = glob.glob(os.path.join(day,'candidate_data','*.dat'))
            numCands = len(candList) ## tp compare number of fit candidates to number of submitted
            fitCands.remove(os.path.join(day,'candidate_data/')) ## might be unnecessary
        else:
            numCands = len(fitCands)
        
        ## count number of fits completed for each model
        numFits = {model: len(os.path.join(day,'*',model+'_result.json')) for model in models}
        sumFits = sum(numFits.values())
        ## count number of candidates that weren't fit
        numUnfit = numCands - len(fitCands) 

        return {
        'fitCands': fitCands,
        'numCands':numCands, 
        'numFits':numFits, 
        'sumFits':sumFits, 
        'numUnfit':numUnfit
        }
    else:
        print('provide a day to count fits for!\n')
        exit(1) ## irreconciable error, hence exit(1)


def get_dataframe(candDir=args.candDir, fitDir=args.fitDir, models=args.models, save=True, file=None):
    '''
    Creates or loads in a pandas dataframe with relevant values for different candidates. If a file is provided, the dataframe will be loaded from that file. Otherwise, the dataframe will be created from the candidate and fit directories provided. 
    Note: may want to include the current things that are expected in dataframe in this description

    Args:
    candDir: path to candidate directory 
    fitDir: path to fit directory 
    models: list of models to search for/consider
    save: boolean to determine whether to save the dataframe to a file to be accessed later
    file: path of saved dataframe to be read in. If None, will proceed to generate dataframe
    '''
    startTime = time.time() ## for timing purposes

    if file:
        print('loading dataframe from file: {}'.format(file)) if args.verbose else None
        df = pd.read_csv(file,index_col=0).fillna(value=np.nan) ## needs to be tested to ensure compatibility with saved dataframe
        df['startDate'] = pd.to_datetime(df['startDate'])
        df['stopDate'] = pd.to_datetime(df['stopDate'])
        return df ## don't need an else since the function will exit if file is provided
    
    ## need to explicitly add all columns here maybe? Will mess with any additional parameters provided to get_json if that is added to this function in the future
    col = ['day','startDate','stopDate','dayPath','cand','candPath','model', 'fitPath','json','fitBool','sampling_time', 'sampler', 'log_evidence', 'log_evidence_err', 'log_noise_evidence', 'log_bayes_factor'] ## addition of start and stop day needs to be tested
    df = pd.DataFrame(columns=col) ## create empty dataframe with columns
    ## set the type for the columns that will be added to the dataframe
    df['day'] = df['day'].astype('str')
    df['startDate'] = df['startDate'].astype(np.datetime64) ## could use this as a way to set bounds on the data that's collected ahead of time
    df['stopDate'] = df['stopDate'].astype(np.datetime64) ## going to use convention that stopDate is the day to be plotted, as that corresponds to the day of the last observation
    df['dayPath'] = df['dayPath'].astype('str')
    df['cand'] = df['cand'].astype('str')
    df['candPath'] = df['candPath'].astype('str')
    df['model'] = df['model'].astype('str')
    df['fitPath'] = df['fitPath'].astype('str')
    df['json'] = df['json'].astype('str')
    df['fitBool'] = df['fitBool'].astype('bool')

    ## should probably change these to be an additional argument passed to get_json, specifically the params argument
    ## are there other useful parameters in the json file that should be included?

    df['sampling_time'] = df['sampling_time'].astype('float')
    df['sampler'] = df['sampler'].astype('str')
    df['log_evidence'] = df['log_evidence'].astype('float')
    df['log_evidence_err'] = df['log_evidence_err'].astype('float')
    df['log_noise_evidence'] = df['log_noise_evidence'].astype('float')
    df['log_bayes_factor'] = df['log_bayes_factor'].astype('float')
    
    dayPathList = glob.glob(os.path.join(candDir, "*",'')) ## list of paths to the days that have candidates
    print('dayPathList: {}\n'.format(dayPathList)) if args.verbose else None
    dayList = [dayPath.split('/')[-2] for dayPath in dayPathList]

    idx = 0 ## used to keep track of the index of the dataframe when defining new values

    for day, dayPath in zip(dayList, dayPathList):
        ## get lists for day level directories
        candPathList = glob.glob(os.path.join(dayPath, "*.csv")) ## could change to have a .dat argument option
        candList = [cand.split('/')[-1].split('.')[0].split('_')[1] for cand in candPathList] ## this is a bit of a mess, but it works (hopefully)
        for cand, candPath in zip(candList, candPathList): ## works around the issue of candidate_data being present in the candidate_fits directory, which is not the case for the countDailyFits function
            ## search for models at same time as candidate data
            for model in models:
                df.at[idx, 'day'] = day
                startDate, stopDate = df.at[idx, 'day'].split('-', 1) ## create values for start and stop day columns
                startDate, stopDate = Time(startDate, format='jd').datetime64, Time(stopDate, format='jd').datetime64 
                ## might be inefficient, adding the date columns makes sample take 1.02 seconds, without it takes 0.88 seconds (15% increase)
                ## actually, running a second time, it only takes 0.81 seconds with the date columns, so it might not matter
                df.at[idx, 'startDate'] = startDate
                df.at[idx, 'stopDate'] = stopDate
                df.at[idx, 'dayPath'] = dayPath
                df.at[idx, 'cand'] = cand
                df.at[idx, 'candPath'] = candPath
                df.at[idx, 'model'] = model

                ## check if fit was completed
                fitPath = os.path.join(fitDir, day, cand,"")
                print('fitPath: {}'.format(fitPath)) if args.verbose else None
                df.at[idx, 'fitPath'] = fitPath
                ## now find json
                jsonPath = os.path.join(fitPath, model+'_result.json')
                jsonBool = True if os.path.exists(jsonPath) else False
                
                print('jsonPath: {}'.format(jsonPath)) if args.verbose else None
                if jsonBool:
                    df.at[idx, 'json'] = jsonPath
                    df.at[idx, 'fitBool'] = True
                    ## now get values from json
                    jsonDict = get_json(file=jsonPath)
                    for key, value in jsonDict.items():
                        df.at[idx, key] = value
                elif not jsonBool:
                    df.at[idx, 'json'] = np.nan
                    df.at[idx, 'fitBool'] = False
                    ## now get values from json
                    jsonDict = get_json(file=False)
                    for key, value in jsonDict.items():
                        df.at[idx, key] = np.nan ## should be np.nan
                idx += 1
                print('get_dataframe idx: {}'.format(idx)) if args.verbose else None
    
    df.sort_values(by=['day','cand','model'], inplace=True)
    df.to_csv(plotDir(name='statsDataframe',ext='.csv')) if save else None
    ## Not exactly the intended use of plotDir, but it works (probably)
    print('completed dataframe creation') if args.verbose else None
    
    print('time to create dataframe: {} seconds\n'.format(round(time.time()-startTime,2))) if args.verbose else None

    return df ## generally, most items returned in df will be strings, with a small number of bools and np.nan values



## Functions to plot daily candidate stats
## these functions could probably be combined for ease of calling, perhaps with argument to determine which plot(s) to make
def plotCands(df, save=True, outdir=args.outdir, ext='.png'): 
    '''
    plot the number of candidates per day as both a line plot (numDailyCand) and histogram (numDailyCandHist), plot rolling average of number of candidates (numDailyCandRolling),
    plot cumulative number of candidates over time (cumDailyCand)
    
    Args:
    df: dataframe with candidate data from get_dataframe function
    save: boolean to determine whether to save the plot or not
    '''
    startTime = time.time()
    print('starting candidate plotting') if args.verbose else None
    ## create subdirectory for plots
    subdir = os.path.join(outdir,'candidates')
    if not os.path.exists(subdir):
        os.mkdir(subdir)
        
    ## create dataframes for plotting
    ## grouped by day and candidate
    df_c= df.groupby(['startDate','stopDate','cand'],as_index=False).agg(tuple).applymap(lambda x: np.array(x))
    ## grouped by day
    df_cd = df_c.groupby(['startDate','stopDate'],as_index=False).agg(tuple).applymap(lambda x: np.array(x))
    df_cd['numCand'] = [len(cand) for cand in df_cd['cand']]
    #print('df_cd: {}\n'.format(df_cd)) if args.verbose else None
    #[print(cand) for cand in df_cd['numCand']]
    #print('total number of candidates: {}'.format(df_cd['numCand'].sum())) if args.verbose else None
    #print(df_cd)
    #df_cdc.to_csv('./df_cdaily.csv')
    ## unique candidates
    df_u = df.groupby(['startDate','stopDate','cand'],as_index=False).agg(tuple).drop_duplicates(subset=['cand'])
    ## unique candidates grouped by day
    df_ud = df_u.groupby(['startDate','stopDate'],as_index=False).agg(tuple).applymap(lambda x: np.array(x))
    df_ud['numCand'] = [len(cand) for cand in df_ud['cand']]
    #print('total number of unique candidates: {}'.format(df_ud['numCand'].sum())) if args.verbose else None
    
    #print('largest nuber of candidates in a single day: {}'.format(df_cd['numCand'].max())) if args.verbose else None
    
    ## plot number of candidates per day
    fig, ax = plotstyle(figsize=(20,15), facecolor='white') 
    sns.histplot(data=df_cd, x='stopDate', weights='numCand', 
                 bins=df_cd['stopDate'].size*2,
                 ax=ax, color='black')
    plt.xticks(rotation=15)
    ax.set_xlabel("Date") 
    ax.set_ylabel('Candidates Per Day')
    plt.savefig(plotDir("numDailyCand",outdir=subdir,ext=ext)) if save else None
    #print('completed numDailyCand plot') if args.verbose else None
    plt.close()
    
     
    ## plot histogram of number of candidates per day
    fig, ax = plotstyle(figsize=(12,8), facecolor='white')
    sns.histplot(df_cd['numCand'], kde=True, 
                 bins=df_cd['numCand'].max(), ax=ax) ## I think having bins equal to the max number of candidates per day looks best
    ax.set_xlabel("Candidates Per Day")
    ax.set_ylabel('Count')
    plt.savefig(plotDir("numDailyCandHist",outdir=subdir,ext=ext)) if save else None
    #print('completed numDailyCandHist plot') if args.verbose else None
    plt.close()
    
    #plot 7 day rolling average of candidates per day
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    sns.histplot(data=df_cd, x='stopDate', weights=df_cd['numCand'].rolling(7).mean(),
                 bins=df_cd['stopDate'].size*2,
                 color='black',linewidth=2) ## note: this won't work with one week of data
    plt.xticks(rotation=15)
    ax.set_xlabel("Date")
    ax.set_ylabel('Candidates Per Day\n(Rolling Average)') ## needs title
    plt.savefig(plotDir("numDailyCandRolling",outdir=subdir,ext=ext)) if save else None
    #print('completed numDailyCandRolling plot') if args.verbose else None
    plt.close()
    
    ## plot cumulative number of candidates per day
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    sns.lineplot(data=df_cd, x='stopDate', y=df_cd['numCand'].cumsum(),
                 color='black',linewidth=2, ax=ax ) 
    plt.xticks(rotation=15)
    ax.set_xlabel("Date")
    ax.set_ylabel('Candidate Count')
    plt.savefig(plotDir("cumDailyCand",outdir=subdir,ext=ext)) if save else None
    #print('completed cumDailyCand plot') if args.verbose else None
    plt.close()
    
    ## plot number of unique candidates
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    sns.histplot(data=df_cd, x='stopDate', weights='numCand',
                 bins=df_cd['stopDate'].size*2,
                 color='black', alpha=0.5, ax=ax)
    sns.histplot(data=df_ud, x='stopDate', weights='numCand',
                 bins=df_cd['stopDate'].size*2,
                 color='blue', ax=ax)
    plt.xticks(rotation=15)
    ax.set_xlabel("Date")
    ax.set_ylabel('Candidate Count')
    plt.savefig(plotDir("uniqueDailyCand",outdir=subdir,ext=ext)) if save else None
    #print('completed uniqueDailyCand plot') if args.verbose else None
    plt.close()
    
    
    print('completed candidate plotting') if args.verbose else None
    print('time to plot candidates: {} seconds\n'.format(round(time.time()-startTime,2))) if args.verbose else None
    
    

## Functions to plot fitting stats
## need a daily fits plot to be made in addition to the cumulative one
def plotFits(df, models=args.models, save=True, outdir=args.outdir, ext='.png'): 
    '''
    plot the cumulative number of fits for each model
    
    Args:
    df: dataframe with candidate data from get_dataframe function
    models: list of models to search for
    save: boolean to determine whether to save the figure or not
    '''
    startTime = time.time()
    print('starting fit plotting') if args.verbose else None
    ## create subdirectory for plots
    subdir = os.path.join(outdir,'fits')
    if not os.path.exists(subdir):
        os.mkdir(subdir)
    ## modelDict creates dict of cumulative fit counts for each model so they can be plotted together
    modelDict = {}
    ## get count of days and unique dates for plotting
    dayList = df['day'].unique()
    dateIdx = df['day'].drop_duplicates().index
    dateList = df['stopDate'][dateIdx] ## this is the date of the last observations made for the fitting
    #print('dayList: {}'.format(dayList)) if args.verbose else None
    #print('dateList: {}\n'.format(dateList)) if args.verbose else None
    ## number of daily candidates
    
    numDaily = np.array([len(df[df['day'] == day]['candPath'].unique()) for day in dayList])
    cumDaily = np.cumsum(numDaily)
    
    ## create dataframes for plotting
    ## grouped by day and candidate
    df_c= df.groupby(['startDate','stopDate','cand'],as_index=False).agg(tuple).applymap(lambda x: np.array(x))
    ## grouped by day
    df_cdc = df_c.groupby(['startDate','stopDate'],as_index=False).agg(tuple).applymap(lambda x: np.array(x))
    df_cdc['numCand'] = [len(cand) for cand in df_cdc['cand']]
    #df_cdc.to_csv('./df_cdaily.csv')
    ## fit-only dataframe
    df_fc= df[df['fitBool']==True].groupby(['startDate','stopDate','cand'],as_index=False).agg(tuple).applymap(lambda x: np.array(x))
    df_fdc = df_fc.groupby(['startDate','stopDate'],as_index=False).agg(tuple).applymap(lambda x: np.array(x))
    df_fdc['numCand'] = [len(cand) for cand in df_fdc['cand']]
    
    for model in models: ## get cumulative number of fits for each model, plot, save, and add to modelDict
        ## compile cumulative number of fits for each model
        modelCount = np.array([len(df[(df['model']==model) & (df['day'] == day) & (df['fitBool'] == True)]) for day in dayList])
        modelCum = np.array(modelCount.cumsum())
        #print('modelCum: {}'.format(modelCum) if args.verbose else None
        modelDict[model] = modelCum
        #print('modelCum: {}'.format(modelCum)) if args.verbose else None
        
        ## plot cumulative number of fits for each model
        ## perhaps this could be a grid of subplots
        fig, ax = plotstyle(figsize=(20,15), facecolor='white')
        ax.plot(dateList,modelCum, label=model)
        ax.plot(dateList, cumDaily, label='Candidate Count', color='black', linewidth=4)
        ax.set_xlabel("Date")
        plt.xticks(rotation=15)
        ax.set_ylabel('Count')
        ax.set_title('{}'.format(model))
        ax.legend()
        plt.savefig(plotDir("cumDailyFits_"+model,outdir=subdir,ext=ext)) if save else None
        #print('completed cumDailyFits plot for {} \n'.format(model)) if args.verbose else None
        plt.close()
    try: ## using a try here because this could totally break if the modelDict has different lengths for each model
        modelDict['Total'] = sum(map(np.array,modelDict.values())).tolist()
    except:
        print('Keys in modelDict probably do not have the same length')
        pass
    ## now plot all models together
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    for key, value in modelDict.items(): 
        ax.plot(dateList,value, label=key, alpha=0.7, linewidth=4) if key != 'Total' else None## need to make a colormap for better visualization
    ax.plot(dateList, cumDaily, label='Candidate Count', color='black', linewidth=4)
    plt.xticks(rotation=15)
    ax.set_xlabel("Date")
    ax.set_ylabel('Count')
    ## need to cmap or something for controlling colors
    #ax.set_title('Cumulative Number of Fits')
    ax.legend()
    plt.savefig(plotDir("cumDailyFitsAll",outdir=subdir,ext=ext)) if save else None ## need to make a version that adds a residual plot below to compare models
    ax.set_yscale('log')
    plt.savefig(plotDir("cumDailyFitsAllLog",outdir=subdir,ext=ext)) if save else None
    #print('completed cumDailyFitsAll plot\n') if args.verbose else None
    plt.close()
    
    ## plot the relative performance of each model against the others
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    rel_perf = {}
    for key, value in modelDict.items():
        if key == 'Total' or key == 'Piro2021':
            continue
        for key2, value2 in modelDict.items():
            #print('rel_perf: {}'.format(rel_perf)) if args.verbose else None
            if key == key2 or key2 == 'Total' or key2 == 'Piro2021':
                continue
            diff = np.array(value)-np.array(value2)
            print('key: {}, key2: {}'.format(key,key2)) if args.verbose else None
            #print('diff: {}'.format(diff)) if args.verbose else None
            if key2+'-'+key in rel_perf.keys(): ## to prevent duplicates in plot
                continue
            rel_perf[key+'-'+key2] = diff
            ax.plot(dateList,diff, label=key+'-'+key2)
    plt.xticks(rotation=15)
    ax.set_xlabel("Date")
    ax.set_ylabel('Difference in Fit Count')
    ax.legend()
    plt.savefig(plotDir("cumDailyFitsRelPerf",outdir=subdir,ext=ext)) if save else None
    #print('completed cumDailyFitsRelPerf plot\n') if args.verbose else None
    plt.close()
    #print('completed cumDailyFits plot for all models \n') if args.verbose else None
    
    print('completed fit plotting') if args.verbose else None
    print('time to plot candidate fits: {} seconds\n'.format(round(time.time()-startTime,2))) if args.verbose else None
    

def plotUnfit(df, models= args.models, save=True, outdir=args.outdir, ext='.png'): ## assumes use of dataframe
    '''
    Plot the number of candidates that were not fit for each day

    Args:
    df: dataframe containing the stats data (expected to be output of get_dataframe) (required)
    models: list of models to search for
    save: boolean to determine whether to save the figure or not
    '''
    startTime = time.time()
    print('starting unfit plotting') if args.verbose else None
    subdir = os.path.join(outdir,'unfit')
    if not os.path.exists(subdir):
        os.mkdir(subdir)
    ## compiling data for plotting
    ## get count of days and unique dates for plotting
    dayList = df['day'].unique()
    dateIdx = df['day'].drop_duplicates().index
    dateList = df['stopDate'][dateIdx] ## this is the date of the last observations made for the fitting
    #print('dayList: {}'.format(dayList)) if args.verbose else None
    #print('dateList: {}\n'.format(dateList)) if args.verbose else None

    ## find number of candidates that were not fit for each day, seperated by model
    ## df uses conditionals in list comprehension, which is wrapped in a dict comprehension
    ## slightly long expression, but should be efficient (dataframe filtering could be slow potentially)
    unfit = {model: 
    np.array([len(df[(df['fitBool'] == False) & (df['model'] == model) & (df['day'] == day)])
    for day in dayList])
    for model in models}
    unfit['Total'] = np.array([len(df[ (df['fitBool'] == False) & (df['day'] == day)]) for day in dayList])

    ## find number of candidates that were fit for each day, seperated by model (for plotting stats later)
    fit = {model: 
    np.array([len(df[(df['fitBool'] == True) & (df['model'] == model) & (df['day'] == day)])
    for day in dayList])
    for model in models}
    fit['Total'] = np.array([len(df[(df['fitBool'] == True) & (df['day'] == day)]) for day in dayList])

    ## total number of fit and unfit per day (for plotting stats later)
    allfit = {model: 
    np.array([len(df[ (df['model'] == model) & (df['day'] == day)])
    for day in dayList])
    for model in models}
    allfit['Total'] = np.array([len(df[ (df['day'] == day)]) for day in dayList])
    

    ## data plotting
    ## plot the number of candidates that were not fit for each day
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    for key, value in unfit.items(): ## one line conditional here is to exclude the total from the histogram
        ax.plot(dateList, value, label=key, alpha=0.6) if key != 'Total' else None
    ax.set_xlabel("Date")
    ax.set_ylabel('Unfit Models')
    #ax.set_title('Number of Unfit Candidates') ## should these have titles?
    plt.xticks(rotation=15)
    ax.legend()
    plt.savefig(plotDir("numDailyUnfit",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## plot a fraction of how many candidates were fit for each day (fit-unfit)/total
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    for key, value in unfit.items(): 
        fracFit = (fit[key]-value)/allfit[key]
        ax.plot(dateList, fracFit, label=key, alpha=0.8) if key != 'Total' else None
    ax.set_xlabel("Date")
    ax.set_ylabel(r'$\frac{Fit-Unfit}{Total}$')
    #ax.set_title('Number of Unfit Candidates') ## should these have titles?
    plt.xticks(rotation=15)
    ax.legend()
    plt.savefig(plotDir("fracDailyUnfit",outdir=subdir,ext=ext)) if save else None
    plt.close()
        

    ## should fix styling as it's currently unclear
    ## plot histogram of number of candidates that were not fit for each day by model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    for key, value in unfit.items():
        sns.histplot(value, label=key,alpha=0.75, ax=ax) if key != 'Total' else None 
    ax.set_xlabel("Daily Unfit Count")
    ax.set_ylabel('Count')
    #ax.set_title('Number of Unfit Candidates per Day') ## should these have titles?
    ax.legend()
    plt.savefig(plotDir("numDailyUnfitModelHist",outdir=subdir,ext=ext)) if save else None
    plt.close()

    ## plot histogram of number of candidates that were not fit for each day
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    sns.histplot(unfit['Total'], bins=20, ax=ax) ## could fine tune the number of bins
    ax.set_xlabel("Number Unfit")
    ax.set_ylabel('Count')
    #ax.set_title("Number of Unfit Candidates per Day")
    plt.savefig(plotDir("numDailyUnfitTotalHist",outdir=subdir,ext=ext)) if save else None
    plt.close()

    ## plot rolling average of number of candidates that were not fit for each day
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    for key, value in unfit.items():
        ax.plot(dateList, pd.Series(value).rolling(7).mean(), label=key)
    ax.set_xlabel("Date")
    ax.set_ylabel('Unfit models\n (7 day rolling average)')
    #ax.set_title('Number of Unfit Candidates') ## should these have titles?
    plt.xticks(rotation=15)
    ax.legend()
    plt.savefig(plotDir("numDailyUnfitRolling",outdir=subdir,ext=ext)) if save else None
    plt.close()

    ## plot cumulative number of candidates that were not fit for each day
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    for key, value in unfit.items():
        ax.plot(dateList, np.cumsum(value), label=key)
    ax.set_xlabel("Date")
    ax.set_ylabel('Cumulative Unfit')
    #ax.set_title('Cumulative Number of Unfit Candidates') ## should these have titles?
    plt.xticks(rotation=15)
    ax.legend()
    plt.savefig(plotDir("cumDailyUnfit",outdir=subdir,ext=ext)) if save else None
    plt.close()
   
    ## plot fraction of candidates that were not fit for each day
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    for key, value in unfit.items():
        fracValue = value/allfit['Total']
        ax.plot(dateList, fracValue, label=key) if key != 'Total' else None
    ax.set_xlabel("Date")
    ax.set_ylabel('Unfit Ratio')
    #ax.set_title('Fraction of Unfit Candidates to Total') ## should these have titles?
    plt.xticks(rotation=15)
    ax.legend()
    plt.savefig(plotDir("fracDailyUnfit",outdir=subdir,ext=ext)) if save else None
    plt.close()

    ## plot rolling average of fraction of candidates that were not fit for each day
    fig, ax = plt.subplots(figsize=(20,15), facecolor='white')
    for key, value in unfit.items(): ## is this the correct method for rolling average ratio?
        fracValue =  pd.Series(value).rolling(7).mean()/pd.Series(allfit['Total']).rolling(7).mean()
        ax.plot(dateList, fracValue, label=key) if key != 'Total' else None
    ax.set_xlabel("Date")
    plt.xticks(rotation=15)
    ax.set_ylabel('Unfit Ratio')
    #ax.set_title('Fraction of Unfit Candidates to Total \n (One Week Rolling Average)') ## should these have titles?
    ax.legend()
    plt.savefig(plotDir("fracDailyUnfitRolling",outdir=subdir,ext=ext)) if save else None
    plt.close()

    '''
    ## this seems to be busted in some way 
    ## (actually, it might not be the most useful plot)
    ## plot cumulative fraction of candidates that were not fit for each day 
    fig, ax = plt.subplots(figsize=(20,15), facecolor='white')
    for key, value in allfit.items():
        fracValue = np.cumsum(value)/np.cumsum(allfit['Total'])
        ax.plot(dayCount, np.cumsum(fracValue), label=key) if key != 'Total' else None
    ax.set_xlabel("Days Since Start")
    ax.set_ylabel('Unfit Ratio\n (Cumulative)')
    #ax.set_title('Cumulative Fraction of Unfit Candidates to Total') ## should these have titles?
    ax.legend()
    plt.savefig(plotDir("cumFracDailyUnfit")) if save else None
    plt.close()
    
    ## plot rolling average of cumulative fraction of candidates that were not fit for each day
    fig, ax = plt.subplots(figsize=(20,15), facecolor='white')
    for key, value in allfit.items():
        fracValue = pd.Series(np.cumsum(value)).rolling(7).mean()/pd.Series(np.cumsum(allfit['Total'])).rolling(7).mean()
        ax.plot(dayCount, fracValue, label=key) if key != 'Total' else None
    ax.set_xlabel("Days Since Start")
    ax.set_ylabel('Ratio')
    #ax.set_title('Cumulative Fraction of Unfit Candidates to Total \n (One Week Rolling Average)') ## should these have titles?
    ax.legend()
    plt.savefig(plotDir("cumFracDailyUnfitRolling")) if save else None
    plt.close()
    '''
    
    print('completed unfit plotting') if args.verbose else None
    print('time to plot unfit: {} seconds\n'.format(round(time.time()-startTime,2))) if args.verbose else None
    
    
    ## maybe a simple bar chart of unfit candidates? (could be useful for a quick glance)
    ## could also add a stacked bar chart of fit and unfit candidates

    ## probably don't need to return anything here



## plot the fit time statistics
def plotSamplingTimes(df, models=args.models, save=True, outdir=args.outdir, ext='.png'):
    '''
    Plot the sampling time statistics for the given dataframe.

    Args:
    df: dataframe containing the stats data (expected to be output of get_dataframe) (required)
    models: list of models to search for
    save: boolean to determine whether to save the figure or not
    '''
    startTime = time.time()
    print('plotting sampling times') if args.verbose else None
    ## create subdirectory for plots
    subdir = os.path.join(outdir,'samplingTimes')
    if not os.path.exists(subdir):
        os.mkdir(subdir)
    ## get count of days and unique dates for plotting
    # dayList = df['day'].unique()
    # dateIdx = df['day'].drop_duplicates().index
    # dateList = df['stopDate'][dateIdx] ## this is the date of the last observations made for the fitting
    #print('dayList: {}'.format(dayList)) if args.verbose else None
    #print('dateList: {}\n'.format(dateList)) if args.verbose else None
    
    ## group by model and day
    df_f = df[df['fitBool']==True].groupby(['startDate','stopDate','model'],as_index=False).agg(tuple).applymap(lambda x: np.array(x))
    df_f['sampling_time'] = df_f['sampling_time'].apply(lambda x: x[~np.isnan(x)])
    df_f['sampling_time_avg'] = [np.mean(timeset) for timeset in df_f['sampling_time'].to_numpy()]
    df_f['sampling_time_median'] = [np.median(timeset) for timeset in df_f['sampling_time']]
    df_f['sampling_time_total'] = df_f.sampling_time.map(sum) #[np.sum(timeset.flatten()) for timeset in df_f['sampling_time']]
    #df_f['sampling_time_cum_total'] = 0 ## initialize
    # for model in models:
    #     df_f['sampling_time_cum_total'][df_f['model']==model] = np.cumsum(df_f['sampling_time_total'][df_f['model']==model])
    
    # print('total sampling time for all fits: {} seconds'.format(round(df_f['sampling_time_total'].sum(),3))) if args.verbose else None
    # df_f.to_csv('test.csv')
    #print(df_f['sampling_time_avg'])
    
    ## group by day
    df_fd = df[df['fitBool']==True].groupby(['startDate','stopDate'],as_index=False).agg(tuple).applymap(lambda x: np.array(x))
    df_fd['sampling_time'] = df_fd['sampling_time'].apply(lambda x: x[~np.isnan(x)])
    df_fd['numCand'] = [len(cand) for cand in df_fd['cand']]
    df_fd['sampling_time_avg'] = [np.mean(timeset) for timeset in df_fd['sampling_time']]
    df_fd['sampling_time_median'] = [np.median(timeset) for timeset in df_fd['sampling_time']]
    df_fd['sampling_time_total'] = df_fd.sampling_time.map(sum)
    #print('total sampling time for all fits: {} seconds'.format(round(df_fd['sampling_time_total'].sum(),3))) if args.verbose else None
    #df_fd.to_csv('test1.csv')

    ## data plotting
    ## plot histogram of fit times for each model
    fig, ax = plt.subplots(figsize=(20,15), facecolor='white')
    plot = sns.histplot(data=df, x='sampling_time', hue='model',hue_order=models, 
                legend='full', ax=ax, alpha=0.5)    
    ax.set_xlabel("Sampling Times (s)")
    ax.set_ylabel('Count')
    #ax.set_yscale('log')
    #ax.set_title('Sampling Times for Each Model') ## should these have titles?
    sns.kdeplot(data=df[df['fitBool']==True], x='sampling_time', hue='model',hue_order=models, 
                ax=plot, alpha=0.5)
    #ax.legend(labels= [model for model in df['model'].unique()], title='Model')
    
    plt.savefig(plotDir("fitTimeHistModel",outdir=subdir,ext=ext)) if save else None
    plt.close()


    ## plot histogram of total daily fit time
    # totalDailyFitTime = np.concatenate(fitTime['Total'],axis=None).ravel()
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    #sns.histplot(totalDailyFitTime,ax=ax) ## could fine tune the number of bins
    plot = sns.histplot(data=df, x='sampling_time', hue='model',hue_order=models, 
                multiple='stack',legend='full', ax=ax, alpha=0.5)
    ax.set_xlabel("Sampling Times (s)")
    ax.set_ylabel('Count')
    #ax.set_title('Daily Sampling Times')
    plt.savefig(plotDir("fitTimeHistStack",outdir=subdir,ext=ext)) if save else None
    plt.close()

    ## plot the daily average fit time for each model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    # for key, value in fitTime.items(): 
    #     meanFitTime = [np.mean(fitDay) for fitDay in value]
    #     ax.plot(dateList, meanFitTime, label=key) ## should be right axis?

    plot= sns.histplot(data=df_f, x='startDate',weights='sampling_time_avg', hue='model', hue_order=models,
                multiple='layer',legend='full',ax=ax, alpha=0.5,bins=69)
    ax.set_xlabel("Date")
    plt.xticks(rotation=15)
    ax.set_ylabel('Sampling Time (s)')

    #ax.set_title('Average Daily Sampling Time') ## should these have titles?
    #ax.legend()
    plt.savefig(plotDir("dailyFitTimeAvg",outdir=subdir,ext=ext)) if save else None
    ## could do a version with std error bars as well

    ## plot the daily median fit time for each model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot= sns.histplot(data=df_f, x='startDate',weights='sampling_time_median', hue='model', hue_order=models,
                multiple='layer',legend='full',ax=ax, alpha=1,bins=69)
    ax.set_xlabel("Date")
    plt.xticks(rotation=15)
    ax.set_ylabel('Sampling Time (s)')
    #ax.set_title('Median Daily Sampling Time') ## should these have titles?
    #ax.legend()
    plt.savefig(plotDir("dailyFitTimeMedian",outdir=subdir,ext=ext)) if save else None
    plt.close()

    ## plot the daily mean fit time for each model (rolling average)
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot= sns.histplot(data=df_f, x='startDate',
                       weights=df_f['sampling_time_avg'].rolling(7).mean(), hue='model', hue_order=models,
                       multiple='layer',legend='full',ax=ax, alpha=0.5,bins=69)
    # sns.histplot(data=df_f, x='startDate',y=df_f['sampling_time_avg'].rolling(7).mean(), hue='model', 
    #             legend='full', ax=ax, alpha=0.5)
    ax.set_xlabel("Date")
    plt.xticks(rotation=15)
    ax.set_ylabel('Sampling Time (s)')
    #ax.set_title('Median Daily Sampling Time \n (One Week Rolling Average)') ## should these have titles?
    #ax.legend()
    plt.savefig(plotDir("dailyFitTimeMeanRolling",outdir=subdir,ext=ext)) if save else None
    plt.close()

    ## plot the cumulative daily fit time for each model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot= sns.histplot(data=df_f, x='startDate',weights='sampling_time_total', hue='model', hue_order=models,
                       multiple='layer',legend='full', cumulative=True,
                       ax=ax, alpha=0.5,bins=69)
    ax.set_xlabel("Date")
    plt.xticks(rotation=15)
    ax.set_ylabel('Sampling Time (s)')
    #ax.set_title('Cumulative Sampling Time') ## should these have titles?
    #ax.legend()
    plt.savefig(plotDir("cumFitTime",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## plot the cumulative daily fit time for each model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    
    plot= sns.histplot(data=df_f, x=df_f['stopDate'],weights=df_f['sampling_time_total']/3600, hue='model', hue_order=models,
                       multiple='stack', #labels=['Kilonova', 'GRB Afterglow', 'Supernova'],#legend='full', 
                       cumulative=True,
                       ax=ax, alpha=0.75,bins=69,linewidth=2)
    ax.set_xlabel("Date",fontsize='large')
    plt.xticks(rotation=15)
    ax.set_ylabel('Sampling Time (hours)',fontsize='large')
    #ax.set_title('Cumulative Sampling Time') ## should these have titles?
    #ax.legend()
    plt.savefig(plotDir("cumFitTimeStack",outdir=subdir,ext=ext)) if save else None
    ax2 = plt.twinx()
    sns.lineplot(data=df_fd, x='stopDate', y=df_fd['numCand'].cumsum()/2,
                 color='black',linewidth=4, ax=ax2 ) ## not sure why it's double counting, dividing by 2 is a quick fix
    # fix the axis tick allignment issues
    nticks = 6
    ax.yaxis.set_major_locator(ticker.LinearLocator(nticks))
    ax2.yaxis.set_major_locator(ticker.LinearLocator(nticks))
    ax.set_ylim(0,5000)
    ax2.set_ylim(0,2500)
    ax2.set_ylabel('Cumulative Candidates',fontsize='large',rotation=270,labelpad=30)
    ## manually added time spans for maintenance and issues
    ax.axvspan('2021-12-13', '2022-01-14', alpha=0.25, color='black',zorder=10,label='ZTF Maintenance')
    ax.axvspan('2022-01-25', '2022-02-10', alpha=0.25, color='black',zorder=10)
    ax.axvspan("2022-03-13","2022-03-25", alpha=0.25, color='black',zorder=10)
    # ax.axvspan("2022-06-15","2022-06-16", alpha=0.25, color='black',zorder=0)
    #ax.axvspan("2022-09-09","2022-09-14", alpha=0.25, color='black',zorder=10)
    ax.axvspan("2022-08-06","2022-09-16", alpha=0.25, color='black',zorder=10, label='Schoty Issues')
    #plot.legend(labels=['Kilonova', 'GRB Afterglow', 'Supernova'])
    plot.legend_.set_title(None)
    plt.savefig(plotDir("cumFitTimeStackWithCands",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## ecdf plot of sampling times by model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.ecdfplot(data=df_f, x='sampling_time_total', hue='model', hue_order=models,
                        #legend='full',
                        linewidth=4, ax=ax, #labels=['Kilonova', 'GRB Afterglow', 'Supernova']
                        )
    ax.set_xlabel("Sampling Time (s)",fontsize='large')
    ax.set_ylabel('Cumulative Fraction',fontsize='large')
    #plot.legend(labels=['Kilonova', 'GRB Afterglow', 'Supernova'])
    plot.legend_.set_title(None)
    plt.savefig(plotDir("samplingTimeDistModel",outdir=subdir,ext=ext)) if save else None
    ax.set_xscale('log')
    plt.savefig(plotDir("samplingTimeDistModelLog",outdir=subdir,ext=ext)) if save else None
    
    ## violin plot of sampling times for different models
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.violinplot(data=df, x='sampling_time',y='model',hue='model', hue_order=models,
                        split=False,
                          legend=True, cut=0,
                          ax=ax)
    ax.set_xlabel("Sampling Time (s)")
    ax.set_ylabel('')
    ax.legend([],[], frameon=False)
    #ax.legend()
    #ax.set_xlim(right=35000)## presumes a certain max sampling time
    #ax.set_ylim(bottom=-500)
    #sns.move_legend(plot, 'lower right')
    plt.savefig(plotDir("SamplingTimeViolin",outdir=subdir,ext=ext)) if save else None
    ax.set_xscale('log')
    plt.savefig(plotDir("SamplingTimeViolinLog",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## box plot of sampling times for different models
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.boxplot(data=df, x='sampling_time',y='model', hue_order=models,
                          hue='model', 
                          ax=ax)
    ax.set_xlabel("Sampling Time (s)")
    #ax.set_ylabel('Model')
    ax.set_ylabel('')
    ax.legend([],[], frameon=False)
    #ax.legend()
    #ax.set_xlim(right=35000)## presumes a certain max sampling time
    #ax.set_ylim(bottom=-500)
    #sns.move_legend(plot, 'lower right')
    plt.savefig(plotDir("SamplingTimeBox",outdir=subdir,ext=ext)) if save else None
    ax.set_xscale('log')
    plt.savefig(plotDir("SamplingTimeBoxLog",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    
    print('completed sampling times plotting') if args.verbose else None
    print('time to plot sampling times: {} seconds\n'.format(round(time.time()-startTime,2))) if args.verbose else None
 
def plotLikelihood(df, models=args.models, save=True, outdir=args.outdir, ext='.png'):
    '''
    Plot the log_evidence and log_bayes_factor statistics for the given dataframe.

    Args:
    df: dataframe containing the stats data (expected to be output of get_dataframe) (required)
    models: list of models to search for
    save: boolean to determine whether to save the figure or not
    '''
    startTime = time.time()
    print('plotting likelihoods') if args.verbose else None
    ## create subdirectory for plots
    subdir = os.path.join(outdir,'likelihood')
    if not os.path.exists(subdir):
        os.mkdir(subdir)
    ## get count of days and unique dates for plotting
    dayList = df['day'].unique()
    dateIdx = df['day'].drop_duplicates().index
    dateList = df['stopDate'][dateIdx] ## this is the date of the last observations made for the fitting
    #print('dayList: {}'.format(dayList)) if args.verbose else None
    #print('dateList: {}\n'.format(dateList)) if args.verbose else None
    
    ## group by model and day
    df_f = df[df['fitBool']==True].groupby(['startDate','stopDate','model'],as_index=False).agg(tuple).applymap(lambda x: np.array(x))
    
    df_f['sampling_time_avg'] = [np.mean(timeset) for timeset in df_f['sampling_time'].to_numpy()]
    df_f['sampling_time_median'] = [np.median(timeset) for timeset in df_f['sampling_time']]
    #df_f.to_csv('test.csv')
    #print(df_f['sampling_time_avg'])
    
    ## find the best fit for each object and day
    df_fo = pd.DataFrame()
    df_fo_filtered = pd.DataFrame() ## only accept if diff in max and min is more than 1 
    for cand in df['cand'].unique():
        df_cand = df[(df['fitBool']==True) & (df['cand']==cand)]
        for day in df_cand['day'].unique():
            df_cd = df_cand[df_cand['day']==day]
            df_cd_max = df_cd[df_cd['log_bayes_factor']==df_cd['log_bayes_factor'].max()]
            df_cd_min = df_cd[df_cd['log_bayes_factor']==df_cd['log_bayes_factor'].min()]
            df_cd_max['model_worst'] = df_cd_min['model'].to_numpy()[0] ## bad practice but works?
            df_fo = df_fo.append(df_cd_max, ignore_index=True)
            df_cd_diff = df_cd_max['log_bayes_factor'].to_numpy()[0] - df_cd_min['log_bayes_factor'].to_numpy()[0]
            df_cd_diff_norm = np.abs(df_cd_diff / df_cd_max['log_bayes_factor'].to_numpy()[0])
            df_cd_max['log_bayes_factor_diff'] = df_cd_diff
            df_cd_max['log_bayes_factor_diff_norm'] = df_cd_diff_norm
            if df_cd_diff > 8:
                df_fo_filtered = df_fo_filtered.append(df_cd_max, ignore_index=True)
    # df_fo_filtered.to_csv('./msiStats/test_fo_filtered_8.csv')
    # exit()
            
    ## print stats about number of cands best fit for each model
    for model in models:
        print('number of cands best fit by {}: {}'.format(model,len(df_fo[df_fo['model']==model]))) if args.verbose else None
    for model in models:
        print('number of cands best fit by {} (>=-100 BF): {}'.format(model,len(df_fo[df_fo['model']==model][df_fo['log_bayes_factor']>=-100]))) if args.verbose else None
    for model in models:
        print('number of cands best fit by {} (>=-50 BF): {}'.format(model,len(df_fo[df_fo['model']==model][df_fo['log_bayes_factor']>=-50]))) if args.verbose else None
    for model in models:
        print('number of cands best fit by {} (>=-10 BF): {}'.format(model,len(df_fo[df_fo['model']==model][df_fo['log_bayes_factor']>=-10]))) if args.verbose else None
    #df_fo.to_csv('./msiStats/test_fo.csv')
    
    
    ## group by day
    df_fd = df[df['fitBool']==True].groupby(['startDate','stopDate'],as_index=False).agg(tuple).applymap(lambda x: np.array(x))
    df_fd['sampling_time_avg'] = [np.mean(timeset) for timeset in df_fd['sampling_time']]
    df_fd['sampling_time_median'] = [np.median(timeset) for timeset in df_fd['sampling_time']]
    #df_fd.to_csv('test1.csv')
    
    
    ## data plotting
    ## list of marker types for plotting
    marker = itertools.cycle(('^','8','s','p','*','h','H','+','x','D','d','.')) 
    ## plot a histogram of the log evidence for each model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.histplot(data=df, x='log_evidence', hue='model', kde=True, hue_order=models,
                        multiple='layer',legend='full',
                        ax=ax, alpha=0.5)    
    ax.set_xlabel("Log Evidence")
    ax.set_ylabel('Count')
    #ax.legend()
    plt.savefig(plotDir("LogEvidenceHistModel",outdir=subdir,ext=ext)) if save else None
    ax.set_yscale('log')
    plt.savefig(plotDir("LogEvidenceHistModelLog",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## plot a histogram of the log evidence in total
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.histplot(data=df, x='log_evidence', hue='model', kde=True, hue_order=models,
                        multiple='stack',legend='full',
                        ax=ax, alpha=1)
    # for line in ax.lines:
    #     line.set_color('black')      
    ax.set_xlabel("Log Evidence")
    ax.set_ylabel('Count')
    #ax.legend()
    plt.savefig(plotDir("LogEvidenceHistTotal",outdir=subdir,ext=ext)) if save else None
    ax.set_yscale('log')
    plt.savefig(plotDir("LogEvidenceHistTotalLog",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## histogram of the log bayes factor for each model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.histplot(data=df, x='log_bayes_factor', hue='model', kde=True, hue_order=models,
                        multiple='layer',legend='full',
                        ax=ax, alpha=0.5)
    # for line in ax.lines:
    #     line.set_color('black')      
    ax.set_xlabel("Log Bayes Factor")
    ax.set_ylabel('Count')
    #ax.legend()
    plt.savefig(plotDir("LogBayesHistModel",outdir=subdir,ext=ext)) if save else None
    ax.set_yscale('log')
    plt.savefig(plotDir("LogBayesHistModelLog",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## scatter plot of bayes factor vs evidence for each model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.scatterplot(data=df, x='log_bayes_factor',y='log_evidence', hue_order=models,
                       hue='model',
                       legend='full',
                       ax=ax, alpha=0.5)
        
    ax.set_xlabel("Log Bayes Factor")
    ax.set_ylabel('Log Evidence')
    #ax.legend()
    plt.savefig(plotDir("LogBayesEvidenceScatterModel",outdir=subdir,ext=ext)) if save else None
    ax.set_ylim(top=1)
    ax.set_yscale('symlog')
    plt.savefig(plotDir("LogBayesEvidenceScatterModelLog",outdir=subdir,ext=ext)) if save else None
    ax.set_xlim(right=1)
    ax.set_xscale('symlog')
    plt.savefig(plotDir("LogBayesEvidenceScatterModelLogLog",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## scatter plot of sampling time vs bayes factor for each model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.scatterplot(data=df, x='sampling_time',y='log_bayes_factor', hue_order=models,
                       hue='model', style='model',
                       legend=True,
                       ax=ax, alpha=0.3)
    ax.set_xlabel("Sampling Time (s)")
    ax.set_ylabel('Log Bayes Factor')
    #ax.legend()
    plt.savefig(plotDir("SamplingTimeBayesScatterModel",outdir=subdir,ext=ext)) if save else None
    ax.set_ylim(top=1)
    ax.set_yscale('symlog')
    plt.savefig(plotDir("SamplingTimeBayesScatterModelLog",outdir=subdir,ext=ext)) if save else None
    ax.set_xscale('symlog')
    ax.set_xlim(left=0)
    plt.savefig(plotDir("SamplingTimeBayesScatterModelLogLog",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## kde plot of sampling time vs bayes factor for each model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.kdeplot(data=df, x='sampling_time',y='log_bayes_factor', 
                       hue='model', hue_order=models, fill=False,
                       legend='full', #clip=((0,25000),(-500,0)),
                       ax=ax, alpha=1)

    #ax.legend()
    ax.set_xlim(right=25000)## presumes a certain max sampling time
    ax.set_ylim(bottom=-500)
    sns.move_legend(plot, 'lower right')
    plt.savefig(plotDir("SamplingTimeBayesKDEModel",outdir=subdir,ext=ext)) if save else None
    # ax.set_xscale('log')
    # ax.set_yscale('symlog')
    # plt.savefig(plotDir("SamplingTimeBayesKDEModelLog",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## kde plot of sampling time vs bayes factor for each model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.kdeplot(data=df, x='sampling_time',y='log_bayes_factor', 
                       hue='model', hue_order=models, fill=True, #labels=['Kilonova', 'GRB Afterglow', 'Supernova'],
                       #legend='full', #clip=((0,25000),(-500,0)),
                       ax=ax, alpha=0.75)

    ax.set_xlabel("Sampling Time (s)", fontsize='large')
    ax.set_ylabel('Log Bayes Factor', fontsize='large')
    #ax.legend()
    ax.set_xlim(right=25000)## presumes a certain max sampling time
    ax.set_ylim(bottom=-500)
    sns.move_legend(plot, 'lower right')
    plot.legend_.set_title(None)
    #plot.legend(labels=['Kilonova', 'GRB Afterglow', 'Supernova'])
    plt.savefig(plotDir("SamplingTimeBayesKDEModelFill",outdir=subdir,ext=ext)) if save else None
    # ax.set_xscale('log')
    # ax.set_yscale('symlog')
    # plt.savefig(plotDir("SamplingTimeBayesKDEModelLog",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## violin plot of bayes factor for different models
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.violinplot(data=df, x='log_bayes_factor',y='model',
                          hue='model', split=False,
                          cut=0,
                          ax=ax)
    ax.set_xlabel("Log Bayes Factor")
    #ax.set_ylabel('Model')
    ax.set_ylabel('')
    #ax.legend()
    #ax.set_xlim(right=35000)## presumes a certain max sampling time
    #ax.set_ylim(bottom=-500)
    #sns.move_legend(plot, 'lower right')
    ax.legend([],[], frameon=False) ## to remove legend (somewhat unecessary for this plot)
    plt.savefig(plotDir("BayesViolin",outdir=subdir,ext=ext)) if save else None
    ax.set_xscale('symlog')
    ax.set_xlim(right=1)
    plt.savefig(plotDir("BayesViolinLog",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## box plot of sampling times for different models
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.boxplot(data=df, x='log_bayes_factor',y='model', 
                          hue='model', hue_order=models, 
                          ax=ax)
    ax.set_xlabel("Log Bayes Factor")
    #ax.set_ylabel('Model')
    ax.set_ylabel('')
    #ax.legend()
    #ax.set_xlim(right=35000)## presumes a certain max sampling time
    #ax.set_ylim(bottom=-500)
    #sns.move_legend(plot, 'lower right')
    ax.legend([],[], frameon=False)
    plt.savefig(plotDir("BayesBox",outdir=subdir,ext=ext)) if save else None
    ax.set_xscale('symlog')
    ax.set_xlim(right=1)
    plt.savefig(plotDir("BayesBoxLog",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    ## histogram of best likelihoods for each model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.histplot(data=df_fo, x='log_bayes_factor', hue='model', kde=True, hue_order=models,
                        multiple='layer',legend='full',
                        ax=ax, alpha=0.6)
    # for line in ax.lines:
    #     line.set_color('black')      
    ax.set_xlabel("Log Bayes Factor")
    ax.set_ylabel('Count')
    #ax.legend()
    plt.savefig(plotDir("LogBayesBestHist",outdir=subdir,ext=ext)) if save else None
    ax.set_xlim(left=-200, right=10)
    plt.savefig(plotDir("LogBayesBestHistCutoff",outdir=subdir,ext=ext)) if save else None
    plt.close() 
    
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.histplot(data=df_fo, x='log_bayes_factor', hue='model', kde=True, hue_order=models,
                        multiple='stack',legend='full',
                        ax=ax, alpha=0.6)
    # for line in ax.lines:
    #     line.set_color('black')      
    ax.set_xlabel("Log Bayes Factor")
    ax.set_ylabel('Count')
    #ax.legend()
    plt.savefig(plotDir("LogBayesBestHistStack",outdir=subdir,ext=ext)) if save else None
    ax.set_xlim(left=-200, right=10)
    plt.savefig(plotDir("LogBayesBestHistCutoffStack",outdir=subdir,ext=ext)) if save else None
    plt.close() 
    
    
    ## same but filtered to only include models with a BF difference > 8
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.histplot(data=df_fo_filtered, x='log_bayes_factor', hue='model', kde=True, hue_order=models,
                        multiple='layer',legend='full',
                        ax=ax, alpha=0.6)
    # for line in ax.lines:
    #     line.set_color('black')      
    ax.set_xlabel("Log Bayes Factor")
    ax.set_ylabel('Count')
    #ax.legend()
    plt.savefig(plotDir("LogBayesBestFilteredHist",outdir=subdir,ext=ext)) if save else None
    ax.set_xlim(left=-200, right=10)
    plt.savefig(plotDir("LogBayesBestFilteredHistCutoff",outdir=subdir,ext=ext)) if save else None
    plt.close() 
    
    
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.histplot(data=df_fo_filtered, x='log_bayes_factor', hue='model', kde=True, hue_order=models,
                        multiple='stack',legend='full',
                        ax=ax, alpha=0.6)
    # for line in ax.lines:
    #     line.set_color('black')      
    ax.set_xlabel("Log Bayes Factor")
    ax.set_ylabel('Count')
    #ax.legend()
    plt.savefig(plotDir("LogBayesBestFilteredHistStack",outdir=subdir,ext=ext)) if save else None
    ax.set_xlim(left=-200, right=10)
    plt.savefig(plotDir("LogBayesBestFilteredHistCutoffStack",outdir=subdir,ext=ext)) if save else None
    plt.close() 
    
    ## relationship between best fit and difference between best fit and worst fit
    ## want something like this: https://stackoverflow.com/questions/48910486/scatterplot-in-python-make-color-and-shape-according-to-different-categorical-v
    ## the shape would be the worst fit model, the color would be the best fit model
    fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    plot = sns.scatterplot(data=df_fo_filtered, x='log_bayes_factor', y='log_bayes_factor_diff', 
                           hue='model', hue_order=models, s=300,
                           legend='full',
                           ax=ax, alpha=0.75)
    ax.set_xlabel("Log Bayes Factor")
    ax.set_ylabel('Difference in Log Bayes Factor')
    plt.savefig(plotDir("LogBayesBestDiff",outdir=subdir,ext=ext)) if save else None
    plt.close()
    
    # fig, ax = plotstyle(figsize=(20,15), facecolor='white')
    # plot = sns.scatterplot(data=df_fo_filtered, x='log_bayes_factor', y='log_bayes_factor_diff',
    #                   hue='model', hue_order=models, legend='full',
    #                   style='model_worst', style_order=models, markers=['o','+','x'],
    #                   s=300, 
    #                   ax=ax, alpha=0.75)
    # ax.set_xlabel("Log Bayes Factor")
    # ax.set_ylabel('Difference in Log Bayes Factor')
    # plt.savefig(plotDir("LogBayesBestDiffBrotherhood",outdir=subdir,ext=ext)) if save else None
    # plt.close()
    
    
    print('completed likelihood plotting') if args.verbose else None
    print('time to plot candidate likelihoods: {} seconds\n'.format(round(time.time()-startTime,2))) if args.verbose else None
        



## To Do:

## plot rolling average of each model fit time -- maybe unneeded?

## add a file size counter and plotter potentially (would use os.path.filesize())

## perhaps a function that finds the model with the highest log_likelihood for each candidate and then plots some stuff about which models were 'most likely' over time and compared to one another



## command:
## python3 ./stats.py -c ./candidate_data/pipelineStructureExample/candidates/partnership -f  ./candidate_data/pipelineStructureExample/candidate_fits -o ./msiStats  --datafile ./msiStats/statsDataframe.csv --verbose -m Bu2019lm TrPi2018 nugent-hyper Piro2021

df = get_dataframe(candDir=args.candDir, models=args.models, save=False, file=args.datafile)
df = df.copy()[df['model'].isin(args.models)] ## so we only consider the models we want to plot
df_daily = df.groupby(['day','cand'],as_index=False).agg(tuple)

print('date range: {} to {}'.format(df['startDate'].min(),df['stopDate'].max())) if args.verbose else None

## drop_duplicates(subset=['day','cand']) is used to remove duplicate candidates from the same day (ie remove extra model counts) and grouping just by cand will give unique candidates per day
print('average number of candidates per day: {}'.format(round(df.drop_duplicates(subset=['day','cand']).groupby(['day'],as_index=False).count()['cand'].mean(),2))) if args.verbose else None

print('average number of unique candidates per day: {}'.format(df.drop_duplicates(subset=['cand']).groupby('day',as_index=False).count()['cand'].mean())) if args.verbose else None

print('median number of candidates per day: {}'.format(round(df.drop_duplicates(subset=['day','cand']).groupby('day',as_index=False).count()['cand'].median(),0))) if args.verbose else None

print('median number of unique candidates per day: {}'.format(round(df_daily.drop_duplicates(subset=['cand']).groupby('day',as_index=False).count()['cand'].median(),0))) if args.verbose else None

print('total number of daily candidates: {}'.format(len(df_daily['cand']))) if args.verbose else None

print('total number of unique candidates: {}'.format(len(df_daily.drop_duplicates(subset=['cand'])))) if args.verbose else None

print("fraction of unique candidates: {}".format(round(len(df_daily.drop_duplicates(subset=['cand']))/len(df_daily['cand']),2))) if args.verbose else None

print('largest number of candidates in a day: {}'.format(df_daily.groupby('day').count()['cand'].max())) if args.verbose else None

print('largest number of unique candidates in a day: {}'.format(df_daily.drop_duplicates(subset=['cand']).groupby('day').count()['cand'].max())) if args.verbose else None

print('average number of observations per candidate: {}'.format(round(df_daily['cand'].value_counts().mean(),3))) if args.verbose else None

print('median number of observations per candidate: {}'.format(df_daily['cand'].value_counts().median())) if args.verbose else None

print('') if args.verbose else None

print('most observed candidates: \n{}'.format(df_daily['cand'].value_counts().head(5))) if args.verbose else None

print('least observed candidates: \n{}'.format(df_daily['cand'].value_counts(ascending=True).head(5))) if args.verbose else None

print('') if args.verbose else None

## some of these fit fractions might be a bit shaky in terms of calculation (getting too deep in filtering and grouping dataframes)

print('fraction of candidates with at least one fit: {}%'.format(round(df.drop_duplicates(subset=['day','cand']).value_counts(subset=['fitBool'], normalize=True)[1]*100,2))) if args.verbose else None

#print('fraction of candidates succesfully fit to all models: {}%'.format(round(len(df_daily[df['fitBool'].apply(lambda x: False not in x)])/len(df_daily)*100,2))) if args.verbose else None
## implementing above is a bit tricky

for model in args.models:
    print('fit success rate for {}: {}%'.format(model,
                                               round(df[(df['model'] == model)].value_counts(subset=['fitBool'], normalize=True)[1]*100,2))) if args.verbose else None

print('fit success rate overall: {}%'.format(round(df.value_counts(subset=['fitBool'], normalize=True)[1]*100,2))) if args.verbose else None

# print('fraction of fit failures overall: {}'.format(len(df[df['fitBool'] == False])/len(df))) if args.verbose else None

print('') if args.verbose else None

print('total sampling time: {} seconds ({} hours)'.format(round(df['sampling_time'].sum(),1),
                                                          round(df['sampling_time'].sum()/60/60,1))) if args.verbose else None

print('average sampling time per day: {} ({} hours)'.format(round(df['sampling_time'].sum()/len(df['day'].unique()),1),
                                                            round(df['sampling_time'].sum()/len(df['day'].unique())/60/60,1))) if args.verbose else None

print('average sampling time per candidate: {} ({} hours)'.format(round(df['sampling_time'].sum()/len(df['cand'].unique()),1),
                                                                 round(df['sampling_time'].sum()/len(df['cand'].unique())/60/60,1))) if args.verbose else None
## avg per candidate should probably be calculated by finding total sampling time per candidate and dividing by number of unique candidates


# print('sampling time for failed fits: {} seconds ({} hours)'.format(round(df[df['fitBool'] == False]['sampling_time'].sum(),3),
#                                                          round(df[df['fitBool'] == False]['sampling_time'].sum()/60/60,1))) if args.verbose else None

# print('sampling time for succesful fits: {} seconds ({} hours)'.format(round(df[df['fitBool'] == True]['sampling_time'].sum(),3),
#                                                                        round(df[df['fitBool'] == True]['sampling_time'].sum()/60/60,1))) if args.verbose else None
print('') if args.verbose else None

for model in args.models:
    print('total sampling time for {}: {} seconds ({} hours)'.format(model,
                                                                      round(df[df['model'] == model]['sampling_time'].sum(),1),
                                                                      round(df[df['model'] == model]['sampling_time'].sum()/60/60,1))) if args.verbose else None
    print('fraction of total sampling time for {}: {}%'.format(model,
                                                              round(df[df['model'] == model]['sampling_time'].sum()/df['sampling_time'].sum()*100,2))) if args.verbose else None
    print('average sampling time for {}: {} seconds ({} hours)'.format(model,
                                                                      round(df[df['model'] == model]['sampling_time'].mean(),1),
                                                                      round(df[df['model'] == model]['sampling_time'].mean()/60/60,1))) if args.verbose else None
    print('median sampling time for {}: {} seconds ({} hours)'.format(model,
                                                                      round(df[df['model'] == model]['sampling_time'].median(),1),
                                                                      round(df[df['model'] == model]['sampling_time'].median()/60/60,1))) if args.verbose else None
    print('') if args.verbose else None
    
print('') if args.verbose else None
## show average, median, and percentile values of bayes factor for each model
for model in args.models:
    print('bayes factor range for {}: {}, {}'.format(model, 
                                                    round(df[df['model'] == model]['log_bayes_factor'].min(),2),
                                                    round(df[df['model'] == model]['log_bayes_factor'].max(),2))) if args.verbose else None
    print('average bayes factor for {}: {}'.format(model, round(df[df['model'] == model]['log_bayes_factor'].mean(),1))) if args.verbose else None
    print('median bayes factor for {}: {}'.format(model, round(df[df['model'] == model]['log_bayes_factor'].median(),1))) if args.verbose else None
    print('bayes factor 95% percentile for {}: {}, {}'.format(model, 
                                                              round(df[df['model'] == model]['log_bayes_factor'].quantile(0.05),2),
                                                              round(df[df['model'] == model]['log_bayes_factor'].quantile(0.95),2))) if args.verbose else None
    print('') if args.verbose else None



## running functions to plot results
models = ['Kilonova', 'GRB Afterglow', 'Supernova', 'Shock Cooling']
df = get_dataframe(candDir=args.candDir, models=args.models, save=True, file=args.datafile, outdir=args.outdir)
df['model'] = df['model'].replace(args.models, models)
try:
    plotCands(df=df,save=True)
    print('completed daily candidate plots (1)\n') if args.verbose else None
except:
    pass
try:
    plotFits(df=df, save=True)
    print('completed cumulative fit plot (2)\n') if args.verbose else None
except:
    pass
try:
    plotUnfit(df=df, save=True)
    print('completed unfit candidate plot (3)\n') if args.verbose else None
except:
    pass
try:
    plotSamplingTimes(df=df,models=models, save=True)
    print('completed sampling time plot (4)\n') if args.verbose else None
except:
    pass
try:
    plotLikelihood(df=df, models=models, save=True)
    print('completed evidence plot (5)\n') if args.verbose else None
except:
    pass