experiment = Experiment()

# Experiment defaults
experiment.name = 'experiment'
experiment.tags = []
experiment.samples = 0
experiment.model = {'fn': None, 'args': [], 'kwargs': {}}
experiment.optimizer = {'fn': None, 'args': [], 'kwargs': {}}
experiment.sessions = []

# Session defaults
session = AutoMunch()
session.losses = {'solubility': 0, 'l1': 0}
session.seed = random.randint(0, 99)
session.cpus = multiprocessing.cpu_count() - 1
session.device = 'cuda' if torch.cuda.is_available() else 'cpu'
session.log = {'when': []}
session.checkpoint = {'when': []}

# Experiment configuration
for string in args.experiment:
    if '=' in string:
        update = parse_dotted(string)
    else:
        with open(string, 'r') as f:
            update = yaml.safe_load(f)
    # If the current session is defined inside the experiment update the session instead
    if 'session' in update:
        update_rec(session, update.pop('session'))
    update_rec(experiment, update)
parser = ArgumentParser()
parser.add_argument('--model', nargs='+', required=True)
parser.add_argument('--data', nargs='+', required=True, default=[])
parser.add_argument('--options', nargs='+', required=False, default=[])
parser.add_argument('--output', type=str, required=True)

args = parser.parse_args()

# region Collecting phase

# Defaults
model = Munch(fn=None, args=[], kwargs={}, state_dict=None)
data = []
options = AutoMunch()
options.cpus = multiprocessing.cpu_count() - 1
options.device = 'cuda' if torch.cuda.is_available() else 'cpu'
options.output = args.output

# Model from --model args
for string in args.model:
    if '=' in string:
        update = parse_dotted(string)
    else:
        with open(string, 'r') as f:
            update = yaml.safe_load(f)
            # If the yaml file contains an entry with key `model` use that one instead
            if 'model' in update.keys():
                update = update['model']
    update_rec(model, update)