misc.py

"""
@author Jack Ringer
Date: 4/30/2023
Description:
Helper functions for misc tasks
"""
import pandas as pd
import os


def concat_csvs(csv_files: list) -> pd.DataFrame:
    """
    combine the provided csvs into a single dataframe and save
    :param csv_files: list[str], paths to csv files to concat
    :return: pd.DataFrame, the combined csvs
    """
    dfs = []
    for csv_file in csv_files:
        df = pd.read_csv(csv_file)
        dfs.append(df)
    return pd.concat(dfs)


def main():
    dirs = ["all_data/small_radius", "all_data/standard_radius"]
    teleport_forward_dirs = ["all_data/teleport_forward/standard"]
    # for some reason small radii in standard for regular tele forward but not
    # overwhelm
    teleport_forward_over_dirs = ["all_data/teleport_forward/standard",
                                  "all_data/teleport_forward/small_radius"]
    csv_files = ["dark_overwhelm_results.csv", "dark_results.csv",
                 "teleport_forward_overwhelm_results.csv",
                 "teleport_forward_results.csv",
                 "teleport_overwhelm_results.csv", "teleport_results.csv",
                 "vanilla_results.csv"]
    save_dir = "results"
    for csv_file in csv_files:
        paths = []
        chosen_dirs = dirs
        if "teleport_forward_overwhelm" in csv_file:
            chosen_dirs = teleport_forward_over_dirs
        elif "teleport_forward" in csv_file:
            chosen_dirs = teleport_forward_dirs
        for subdir in chosen_dirs:
            paths.append(os.path.join(subdir, csv_file))
        concat_df = concat_csvs(paths)
        save_path = os.path.join(save_dir, csv_file)
        concat_df.to_csv(save_path, index=False)
        print("Saved data for %s to %s" % (csv_file, save_path))


if __name__ == "__main__":
    main()