import os
import numpy as np
import pandas as pd
import argparse

parser = argparse.ArgumentParser(description='PET lymphoma classification')

parser.add_argument('--dir', type=str, default='results', help='name of folder to find convergence files in (default "results")')
parser.add_argument('--metric', type=str, default='auc', help='which metric to compare models with (default "auc")')
args = parser.parse_args()
print(args)

DF = []
files = os.listdir(args.dir)
for f in files: 
    if (f.endswith('.csv')) & (f.startswith('convergence')):
        df = pd.read_csv(os.path.join(args.dir, f))
        df = df[(df.split=='validation')&(df.metric==args.metric)].sort_values(by='epoch')
        if len(df)<7:
            print('SHORT',args.dir,f,'(',len(df),')')
            # continue
        print(args.dir,f,len(df))
        # Get last epoch
        df = df.tail(1)
        df['split'] = int( f.split('_')[1].split('split')[1] )
        df['run'] = int( f.split('run')[1].split('.')[0] )
        DF.append(df)

# Best run in this folder
DF = pd.concat(DF).reset_index(drop=True) 

grouped    = DF.groupby(['split'])
maxidx     = grouped['value'].idxmax()
dfbest     = DF.loc[maxidx]
# Top run in order of split
dfbest = dfbest.sort_values(by='split',ascending=True)
dfbest.to_csv( os.path.join(args.dir,'best_run.csv'),index=False)
print('Best run of ',args.dir,'(',str(len(dfbest)),') :',dfbest.sort_values(by='split',ascending=False).run.values)