import argparse
import subprocess


parser = argparse.ArgumentParser(description="run")
parser.add_argument("--device", type=str, default='0')
parser.add_argument("--run", type=int)

args = parser.parse_args()


if args.run == 0:
    target = 'O2-N2-CO2'
    transition = 'masked'
    diffusion_noise_schedule = 'latent-property-learnable-schedule'
    devices = args.device.split(',')

    exp_name = "polymer_exp"
    subprocess.call(f"export CUDA_VISIBLE_DEVICES={args.device} && python meld/main.py --config-name=config.yaml model.diffusion_noise_schedule={diffusion_noise_schedule} \
    model.ensure_connected=True dataset.task_name={target} dataset.guidance_target={target} general.wandb='disabled' general.exp_name={exp_name} \
    general.gpus={len(devices)} model.transition={transition} general.save_model=True train.batch_size=512 train.n_epochs=15000 \
    model.use_provided_drop_ids=True train.lr=2e-4 model.diffusion_steps=200 model.guide_scale=2 model.transition_hidden_dim=8", shell=True)

