"""
This script automates cleaning up a benchmark/experiment run of some algo
against some config (with possibly more than one tune trial,
e.g. torch=grid_search([True, False])).

Run `python cleanup_experiment.py --help` for more information.

Use on an input directory with trial contents e.g.:
..
IMPALA_BreakoutNoFrameskip-v4_0_use_pytorch=False_2020-05-11_10-17-54topr3h9k
IMPALA_BreakoutNoFrameskip-v4_0_use_pytorch=False_2020-05-11_13-59-35dqaetxnf
IMPALA_BreakoutNoFrameskip-v4_0_use_pytorch=False_2020-05-11_17-21-28tbhedw72
IMPALA_BreakoutNoFrameskip-v4_2_use_pytorch=True_2020-05-11_10-17-54lv20cgn_
IMPALA_BreakoutNoFrameskip-v4_2_use_pytorch=True_2020-05-11_13-59-35kwzhax_y
IMPALA_BreakoutNoFrameskip-v4_2_use_pytorch=True_2020-05-11_17-21-28a5j0s7za

Then run:
>> python cleanup_experiment.py --experiment-dir [parent dir w/ trial sub-dirs]
>>   --output-dir [your out dir] --results-filter dumb_col_2,superfluous_col3
>>   --results-max-size [max results file size in kb before(!) zipping]

The script will create one output sub-dir for each trial and only copy
the configuration and the csv results (filtered and every nth row removed
based on the given args).
"""

import argparse
import json
import os
import re
import shutil
import yaml

parser = argparse.ArgumentParser()
parser.add_argument(
    "--experiment-dir",
    type=str,
    help="Experiment dir in which all sub-runs (seeds) are "
    "located (as sub-dirs). Each sub0-run dir must contain the files: "
    "params.json and progress.csv.",
)
parser.add_argument(
    "--output-dir",
    type=str,
    help="The output dir, in which the cleaned up output will be placed.",
)
parser.add_argument(
    "--results-filter",
    type=str,
    help="comma-separated list of csv fields to exclude.",
    default="experiment_id,pid,hostname,node_ip,trial_id,hist_stats/episode_"
    "reward,hist_stats/episode_lengths,experiment_tag",
)
parser.add_argument(
    "--results-max-size",
    type=int,
    help="the max. size of the final results.csv file (in kb). Will erase "
    "every nth line in the original input to reach that goal. "
    "Use 0 for no limit (default=100).",
    default=100,
)


def process_single_run(in_dir, out_dir):
    exp_dir = os.listdir(in_dir)

    # Make sure trials dir is ok.
    assert (
        "params.json" in exp_dir and "progress.csv" in exp_dir
    ), "params.json or progress.csv not found in {}!".format(in_dir)

    os.makedirs(out_dir, exist_ok=True)

    for file in exp_dir:
        absfile = os.path.join(in_dir, file)
        # Config file -> Convert to yaml and move to output dir.
        if file == "params.json":
            assert os.path.isfile(absfile), "{} not a file!".format(file)
            with open(absfile) as fp:
                contents = json.load(fp)
            with open(os.path.join(out_dir, "config.yaml"), "w") as fp:
                yaml.dump(contents, fp)
        # Progress csv file -> Filter out some columns, cut, and write to
        # output_dir.
        elif file == "progress.csv":
            assert os.path.isfile(absfile), "{} not a file!".format(file)
            col_idx_to_filter = []
            with open(absfile) as fp:
                # Get column names.
                col_names_orig = fp.readline().strip().split(",")
                # Split by comma (abiding to quotes), filter out
                # unwanted columns, then write to disk.
                cols_to_filter = args.results_filter.split(",")
                for i, c in enumerate(col_names_orig):
                    if c in cols_to_filter:
                        col_idx_to_filter.insert(0, i)
                col_names = col_names_orig.copy()
                for idx in col_idx_to_filter:
                    col_names.pop(idx)
                absfile_out = os.path.join(out_dir, "progress.csv")
                with open(absfile_out, "w") as out_fp:
                    print(",".join(col_names), file=out_fp)
                    while True:
                        line = fp.readline().strip()
                        if not line:
                            break
                        line = re.sub(
                            "(,{2,})",
                            lambda m: ",None" * (len(m.group()) - 1) + ",",
                            line,
                        )
                        cols = re.findall('".+?"|[^,]+', line)
                        if len(cols) != len(col_names_orig):
                            continue
                        for idx in col_idx_to_filter:
                            cols.pop(idx)
                        print(",".join(cols), file=out_fp)

            # Reduce the size of the output file if necessary.
            out_size = os.path.getsize(absfile_out)
            max_size = args.results_max_size * 1024
            if 0 < max_size < out_size:
                # Figure out roughly every which line we have to drop.
                ratio = out_size / max_size
                # If ratio > 2.0, we'll have to keep only every nth line.
                if ratio > 2.0:
                    nth = out_size // max_size
                    os.system(
                        "awk 'NR==1||NR%{}==0' {} > {}.new".format(
                            nth, absfile_out, absfile_out
                        )
                    )
                # If ratio < 2.0 (>1.0), we'll have to drop every nth line.
                else:
                    nth = out_size // (out_size - max_size)
                    os.system(
                        "awk 'NR==1||NR%{}!=0' {} > {}.new".format(
                            nth, absfile_out, absfile_out
                        )
                    )
                os.remove(absfile_out)
                os.rename(absfile_out + ".new", absfile_out)

            # Zip progress.csv into results.zip.
            zip_file = os.path.join(out_dir, "results.zip")
            try:
                os.remove(zip_file)
            except FileNotFoundError:
                pass
            os.system(
                "zip -j {} {}".format(zip_file, os.path.join(out_dir, "progress.csv"))
            )
            os.remove(os.path.join(out_dir, "progress.csv"))

        # TBX events file -> Move as is.
        elif re.search("^(events\\.out\\.|params\\.pkl)", file):
            assert os.path.isfile(absfile), "{} not a file!".format(file)
            shutil.copyfile(absfile, os.path.join(out_dir, file))


if __name__ == "__main__":
    args = parser.parse_args()
    exp_dir = os.listdir(args.experiment_dir)
    # Loop through all sub-directories.
    for i, sub_run in enumerate(sorted(exp_dir)):
        abspath = os.path.join(args.experiment_dir, sub_run)
        # This is a seed run.
        if os.path.isdir(abspath) and re.search(
            "^(\\w+?)_(\\w+?-v\\d+)(_\\d+)", sub_run
        ):
            # Create meaningful output dir name:
            # [algo]_[env]_[trial #]_[trial-config]_[date YYYY-MM-DD].
            cleaned_up_out = re.sub(
                "^(\\w+?)_(\\w+?-v\\d+)(_\\d+)(_.+)?(_\\d{4}-\\d{2}-\\d{2})"
                "_\\d{2}-\\d{2}-\\w+",
                "{:02}_\\1_\\2\\4\\5".format(i),
                sub_run,
            )
            # Remove superflous `env=` specifier (anv always included in name).
            cleaned_up_out = re.sub(
                "^(.+)env=\\w+?-v\\d+,?(.+)", "\\1\\2", cleaned_up_out
            )
            out_path = os.path.join(args.output_dir, cleaned_up_out)
            process_single_run(abspath, out_path)
    # Done.
    print("done")