if __name__ == '__main__':
    import argparse
    import os
    from ml_logger import instr
    from model_free_analysis.baselines import RUN  # Use baselines prefix format!
    import jaynes
    from dmc_gen.train import main
    from dmc_gen.config import Args
    from params_proto.neo_hyper import Sweep
    from pathlib import Path

    RUN.prefix = RUN.prefix.replace('{file_stem}', 'baselines/drqv2_original/train.py')

    parser = argparse.ArgumentParser()
    parser.add_argument('sweep_file')
    parser.add_argument('-m', '--mode')
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('-v', '--verbose', action='store_true')
    args = parser.parse_args()

    assert os.path.isfile(args.sweep_file)

    RUN.debug = args.debug
    sweep = Sweep(RUN, Args).load(args.sweep_file)

    jaynes.config(args.mode)
    for kwargs in sweep:
        thunk = instr(main, **kwargs)
        jaynes.run(thunk)

    jaynes.listen()
