import exp_configs
import argparse
#from accelerate import Accelerator
from haven import haven_wizard as hw
import sys

RESULTS_FNAME = "results.ipynb"


def main(exp_dict, savedir, args):

    from openood.pipelines import get_pipeline
    from openood.utils import setup_config

    config_pths = exp_dict['configs']

    # store config files path to system to path them to setup_config with minimal change to original code
    sys.argv = ['main', '--config']
    for config_pth in config_pths:
        sys.argv.append(config_pth)

    config = setup_config()
    pipeline = get_pipeline(config)
    pipeline.run()

    print("Experiment done\n")


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-e",
        "--exp_group",
        help="Define the experiment group to run.",
    )
    parser.add_argument(
        "-sb",
        "--savedir_base",
        required=True,
        help="Define the base directory where the experiments will be saved.",
    )
    parser.add_argument(
        "-r", "--reset", default=0, type=int, help="Reset or resume the experiment."
    )
    parser.add_argument(
        "-j", "--job_scheduler", default=None, help="Choose Job Scheduler."
    )
    parser.add_argument(
        "--python_binary", default="python", help="path to your python executable"
    )
    parser.add_argument(
        "--epochs", default=10, type=int, help="Number of epochs to train."
    )

    args, others = parser.parse_known_args()

    # Choose Job Scheduler
    job_config = None

    if args.job_scheduler == "toolkit":
        import job_configs

        job_config = job_configs.JOB_CONFIG

    # Run experiments and create results file

    hw.run_wizard(
        func=main,
        exp_list=exp_configs.EXP_GROUPS[args.exp_group],
        savedir_base=args.savedir_base,
        reset=args.reset,
        job_config=job_config,
        results_fname=RESULTS_FNAME,
        python_binary_path=args.python_binary,
        args=args,
    )
