import argparse
import os
import re
import subprocess

# import sys
import time
import torch
import requests
import json
from support_alignment.experiment_registry import registry


def multi_gpu_launcher(commands):
    """
    Launch commands on the local machine, using all GPUs in parallel.
    """
    print("WARNING: using experimental multi_gpu_launcher.")
    try:
        # Get list of GPUs from env, split by ',' and remove empty string ''
        # To handle the case when there is one extra comma: `CUDA_VISIBLE_DEVICES=0,1,2,3, python3 ...`
        available_gpus = [
            x for x in os.environ["CUDA_VISIBLE_DEVICES"].split(",") if x != ""
        ]
    except Exception:
        # If the env variable is not set, we use all GPUs
        available_gpus = [str(x) for x in range(torch.cuda.device_count())]
    n_gpus = len(available_gpus)
    procs_by_gpu = [None] * n_gpus

    while len(commands) > 0:
        for idx, gpu_idx in enumerate(available_gpus):
            proc = procs_by_gpu[idx]
            if (proc is None) or (proc.poll() is not None):
                # Nothing is running on this GPU; launch a command.
                cmd = commands.pop(0)
                new_proc = subprocess.Popen(
                    f"CUDA_VISIBLE_DEVICES={gpu_idx} {cmd}", shell=True
                )
                procs_by_gpu[idx] = new_proc
                break
        time.sleep(1)

    # Wait for the last few tasks to finish before returning
    for p in procs_by_gpu:
        if p is not None:
            p.wait()


parser = argparse.ArgumentParser()
parser.add_argument("args", nargs="*")
parser.add_argument("--regex", type=str, required=True, help="experiment name regex")

args = parser.parse_args()

regex = args.regex
print(f"Regex {regex} filter results:")
experiment_list = [
    name for name in registry.list_experiments() if re.match(regex, name)
]
for i, experiment_name in enumerate(experiment_list):
    print(f"{i}: {experiment_name}")
print()

commands = []
for experiment_name in experiment_list:
    # print(f"Starting {experiment_name}")
    command = [
        "python3",
        "-m",
        "support_alignment.scripts.train",
        f"--config_name={experiment_name}",
    ]
    command.extend(args.args or [])
    commands.append(" ".join(command))

for c in commands:
    print(c, flush=True)
multi_gpu_launcher(commands)
