"""
Poor Man's Configurator, just shuttling this code away from train.py
The code in this script then overrides the globals()
"""

from collections import Counter
import os
import sys
from ast import literal_eval


IGNORE = ("local_rank",)
master_process = int(os.environ.get("SLURM_PROCID", os.environ.get("RANK", 0))) == 0

passed_args = [arg.split("=")[0] for arg in sys.argv[1:]]
for arg, count in Counter(passed_args).items():
    if count > 1 and master_process:
        raise ValueError(f"Arg {arg} passed in {count} times from the command line. Likely a bug")

for arg in sys.argv[1:]:
    if "=" not in arg:
        # assume it's the name of a config file
        assert not arg.startswith("--")
        config_file = arg
        if master_process:
            print(f"Overriding config with {config_file}:")
            with open(config_file) as f:
                print(f.read())
        exec(open(config_file).read())
    else:
        # assume it's a --key=value argument
        assert arg.startswith("--")
        key, val = arg.split("=")
        key = key[2:]
        if key in IGNORE:
            print(f"ignoring {key}")
            continue
        if key in globals():
            try:
                # attempt to eval it it (e.g. if bool, number, or etc)
                attempt = literal_eval(val)
            except (SyntaxError, ValueError):
                # if that goes wrong, just use the string
                attempt = val
            # ensure the types match ok
            assert globals()[key] is None or isinstance(attempt, type(globals()[key]))
            # cross fingers
            if master_process:
                print(f"Overriding: {key} = {attempt}")
            globals()[key] = attempt
        else:
            raise ValueError(f"Unknown config key: {key}")