#!/usr/bin/env python3
# File to generate commands to perform experiments

from argparse import ArgumentParser
from yaml import Loader, Dumper, load, dump
from fnmatch import fnmatch
from copy import deepcopy
import itertools
from utils import experiment_slug, saved_measures

from_iterable = itertools.chain.from_iterable

def cartesian_explosion(dct, partial={}):
    """
    Calculates the cartesian product of all choices.
    """
    if len(dct.items()) > 0:
        pname, pvalues = list(dct.items())[0]
        newdct = {p: pv for p, pv in dct.items() if p != pname}
        if not isinstance(pvalues, list):
            pvalues = [pvalues]
        for pval in pvalues:
            newpartial = deepcopy(partial)
            newpartial[pname] = pval
            yield from cartesian_explosion(newdct, partial=newpartial)
    else:
        yield partial


class Attrs:
    def __init__(self, dct):
        self.__dict__.update(dct)


def items_to_attrs(dct):
    "Returns an object with attributes instead of a dictionary"
    return Attrs(dct)

                
if __name__ == "__main__":
    parser = ArgumentParser(description="Generate experiment command line arguments")
    parser.add_argument("expmatch", type=str, default="*", help="Experiments to match, uses fnmatch, i.e. linux-style glob patterns")
    parser.add_argument("--file", type=str, default="experiments.yml", help="Configuration file")
    parser.add_argument("--paramdir", type=str, default="parameters", help="Directory to put parameters files")
    parser.add_argument("--resultdir", type=str, default="./results", help="Directory to put results files")
    parser.add_argument("--no_measure", type=str, nargs="*", default=[], help="Things not to measure")
    parser.add_argument("--prefix", type=str, default="", help="Prefix to assign to the experiment")
    parser.add_argument("--nodvc", action="store_true", default=False, help="Prints non-dvc commands")
    args = parser.parse_args()

    # Prefix of the command to execute
    script = "./run_network.py"

    data = load(open(args.file), Loader=Loader)
    for experiment, ddict in data.items():
        if fnmatch(experiment, args.expmatch):
            print(f"# Experiment {experiment}")
            for current_experiment_params in cartesian_explosion(ddict):
                slug = experiment_slug(items_to_attrs(current_experiment_params))
                cmd_args = [[f"--{pname}", f"{pval}"] for pname, pval in current_experiment_params.items()]
                cmd = list(from_iterable(cmd_args))
                if args.nodvc:
                    print(script, " ".join(cmd))
                else:
                    # 1. Write parameters to param files
                    with open(f"{args.paramdir}/{slug}.yaml", "w") as fout:
                        fout.write(dump(current_experiment_params, default_flow_style=False))
                    # 2. Build the dvc parameter considering dependencies, plots outputs, and param files
                    dvc_name = args.prefix + deepcopy(slug).replace("@", "-").replace(".", "_")
                    plist = ",".join(current_experiment_params.keys())
                    plots = list(from_iterable([["--plots", f"{args.resultdir}/{measure}@{slug}.csv"] for measure in saved_measures if measure not in args.no_measure]))
                    dvcline = ["dvc", "run", "-n", dvc_name,
                               "-d", "run_network.py", "-d", "linear_module.py", "-d", "utils.py",
                               "-p", f"{args.paramdir}/{slug}.yaml:{plist}"] + plots
                    print(" ".join(dvcline), script, " ".join(cmd), "--output_folder", args.resultdir, "--no_measure", " ".join(args.no_measure))
            print("\n")
        else:
            print(f"# SKIPPED {experiment}: NOT MATCHING {args.expmatch}")
