import os
import json

DATA_DIR = "2_stage_vae/"



def make_output_dir(name, overwrite, args, seed_num=None):
    dirname = get_output_dir(name, seed_num)
    if not os.path.exists(dirname):
        os.makedirs(dirname)
    args_name = os.path.join(dirname, "args.json")
    args = vars(args)
    with open(args_name, "w") as f:
        json.dump(args, f)
    return dirname


def get_output_dir(name, seed_num):
    if seed_num is not None:
        dirname = os.path.join(DATA_DIR, name, f"seed_{seed_num}")
    else:
        dirname = os.path.join(DATA_DIR, name)
    return dirname
