#!/usr/bin/env python3
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

"""Runs the train script with a grid of hyperparameters."""
import itertools
import os
import subprocess
from multiprocessing import Process, Queue

# list of GPU IDs that we want to use, one job will be started for every ID in the list
GPUS = [0]
DRY_RUN = False  # if False then print out the commands to be run, if True then run
DATA_DIR = ""  # path to the So2Sat data directory

# Hyperparameter options
model_options = ["resnet18"]
lr_options = [1e-3]
loss_options = ["ce"]
weight_options = ["random"]
seeds = list(range(32))


def do_work(work: "Queue[str]", gpu_idx: int) -> bool:
    """Process for each ID in GPUS."""
    while not work.empty():
        experiment = work.get()
        experiment = experiment.replace("GPU", str(gpu_idx))
        print(experiment)
        if not DRY_RUN:
            subprocess.call(experiment.split(" "))
    return True


if __name__ == "__main__":
    work: "Queue[str]" = Queue()

    for model, lr, loss, weights, seed in itertools.product(
        model_options, lr_options, loss_options, weight_options, seeds
    ):
        experiment_name = f"{model}_{lr}_{loss}_{weights.replace('_','-')}_{seed}"

        output_dir = os.path.join("output", "so2sat_seed_experiments")
        log_dir = os.path.join(output_dir, "logs")
        config_file = os.path.join("conf", "so2sat.yaml")

        if not os.path.exists(os.path.join(output_dir, experiment_name)):
            command = (
                "python train.py"
                + f" config_file={config_file}"
                + f" experiment.name={experiment_name}"
                + f" experiment.module.model={model}"
                + f" experiment.module.learning_rate={lr}"
                + f" experiment.module.loss={loss}"
                + f" experiment.module.weights={weights}"
                + f" experiment.datamodule.weights={weights}"
                + f" program.output_dir={output_dir}"
                + f" program.log_dir={log_dir}"
                + f" program.data_dir={DATA_DIR}"
                + f" program.seed={seed}"
                + " trainer.gpus=[GPU]"
            )
            command = command.strip()

            work.put(command)

    processes = []
    for gpu_idx in GPUS:
        p = Process(target=do_work, args=(work, gpu_idx))
        processes.append(p)
        p.start()
    for p in processes:
        p.join()
