"""
batch_runner: Instead of running individual trials (on different terminals,
if using run scripts), runs a batch of trials with shared resource, and
each trial can be runned in a separate process using multiprocessing.

By default, a batch means the set of trials contained in a run script.
Therefore, this program works by providing a file that contains either
a list of trial paths, or the result that is output by Experiment.generate_run_scripts

Processes have their own memory space, but modern operating systems
implement copy-on-write - If you don't modify the resource, the memory
for it won't be physically copied.
"""
import os
import argparse
import pickle
import multiprocessing

def run_trial(trial, resource, logging):
    trial.set_resource(resource)
    trial.run(logging=logging)

def main():
    parser = argparse.ArgumentParser(description='Run a batch of trials.')
    parser.add_argument("file_path", type=str,
                        help="File that contains a list of trial paths or the output by Experiment.generate_run_scripts")
    parser.add_argument("exp_path", type=str,
                        help="Path to experiment root")
    parser.add_argument("-p", "--num-proc", type=int,
                        help="Number of processes to run trials in parallel",
                        default=4)
    parser.add_argument("--logging", action="store_true")
    parser.add_argument("--spawn", action="store_true")
    args = parser.parse_args()

    if not os.path.exists(args.file_path):
        print("{} not found".format(args.file_path))
        return

    with open(args.file_path) as f:
        lines = f.readlines()

    trial_paths = []
    for line in lines:
        line = line.strip()
        if line.startswith("python trial_runner.py"):
            # This is generated by Experiment.generate_run_scripts
            trial_path = line.split()[2]
            if trial_path.startswith("\""):
                trial_path = trial_path[1:-1]
        else:
            trial_path = line
        trial_paths.append(trial_path)

    trials_to_run = []
    for trial_path in trial_paths:
        # Load the trial
        with open(os.path.join(args.exp_path, trial_path), "rb") as f:
            trial = pickle.load(f)
            if os.path.exists(os.path.join(args.exp_path, trial.name, "config.yaml")):
                print("Skipping {} because it seems to be done".format(trial.name))
            else:
                trials_to_run.append(trial)

    if len(trials_to_run) == 0:
        print("Nothing to run.")
        return

    # Load shared_resource - because it is shared, only one trial needs
    # to provide such a resource.
    resource = None
    for trial in trials_to_run:
        if trial.could_provide_resource():
            resource = trial.provide_shared_resource()
            break
    # No resource is provided. We can still keep going.
    if resource is None:
        print("No resource provided.")
    func_args = [(trial, resource, args.logging)
                 for trial in trials_to_run]
    if args.spawn:
        multiprocessing.set_start_method('spawn')
    with multiprocessing.Pool(processes=args.num_proc) as pool:
        pool.starmap(run_trial, func_args)

if __name__ == "__main__":
    main()
