import glob, os, re
import pandas as pd
import seaborn as sns;
from optparse import OptionParser

sns.set()
import matplotlib.pyplot as plt


def process_data(experimentPath, split):
    results = None
    path = experimentPath + "/" + "**/" + split + "/summary.csv"
    for filename in glob.iglob(path, recursive=True):
        # extract the epoch number
        df = pd.read_csv(filename, index_col=0)
        df['EPOCH'] = int(filename.split("/")[-3:-2][0])

        if results is not None:
            results = pd.concat([results, df], axis=0)

        else:
            results = df
    results['data_split'] = split
    results['expname'] = experimentPath
    return results

if __name__ == "__main__":

    parser = OptionParser()

    parser.add_option("--input", help="specify the input data")

    parser.add_option("--y_columns", help="y column name")
    parser.add_option("--save_dataframe", help="location to save the data frame to ")
    (options, args) = parser.parse_args()

    train = process_data(options.input, "train")
    validation = process_data(options.input, "validation")
    results = pd.concat([train, validation], axis=1)


    results.to_csv(options.save_dataframe)
    ax = sns.lineplot(x="EPOCH", y="LOSS", data=results['train'])
    plt.show()