import json
from pprint import pprint
import subprocess
import random

with open("config_template.json", "r") as f:
    config_template = json.load(f)

variables = [(k, v) for k, v in config_template.items() if isinstance(v, list)]
consts = {k: v for k, v in config_template.items() if not isinstance(v, list)}

# pprint(consts)


def all_assignments(variables, consts):
    if len(variables) == 0:
        yield consts
        return
    k, vs = variables[0]
    for assignment in all_assignments(variables[1:], consts):
        for value in vs:
            if isinstance(value, dict):
                yield {**value, **assignment}
            else:
                yield {k: value, **assignment}


assignments = list(all_assignments(variables, consts))
random.shuffle(assignments)

for assignment in assignments:
    with open("config_hpgs.json", "w") as f:
        json.dump(assignment, f, indent=2)
    subprocess.run(
        [
            "torchrun",
            "--standalone",
            "--nproc-per-node",
            "gpu",
            "distributed.py",
            "config_hpgs.json",
        ]
    )