import os
import tempfile

from monai.bundle.config_parser import ConfigParser
from monai.apps import download_and_extract

from monai.apps.auto3dseg import AutoRunner

if __name__ == '__main__':
    directory = os.environ.get("MONAI_DATA_DIRECTORY")
    root_dir = tempfile.mkdtemp() if directory is None else directory
    print(root_dir)

    msd_task = "Task04_Hippocampus"
    resource = "https://msd-for-monai.s3-us-west-2.amazonaws.com/" + msd_task + ".tar"

    compressed_file = os.path.join(root_dir, msd_task + ".tar")
    dataroot = os.path.join(root_dir, msd_task)
    if not os.path.exists(dataroot):
        download_and_extract(resource, compressed_file, root_dir)

    datalist_file = "build/msd_" + msd_task.lower() + "_folds.json"

    input_cfg = {
        "name": msd_task,  # optional, it is only for your own record
        "task": "segmentation",  # optional, it is only for your own record
        "modality": "MRI",  # required
        "datalist": datalist_file,  # required
        "dataroot": dataroot,  # required
    }
    input = "./input.yaml"
    ConfigParser.export_config_file(input_cfg, input)

    runner = AutoRunner(input=input, hpo=True, ensemble=False)
    num_epoch = 2
    hpo_params = {
        "maxTrialNumber": 20,
        "maxExperimentDuration": "30m",
        "num_epochs_per_validation": 1,
        "num_images_per_batch": 1,
        "num_epochs": 2,
        "num_warmup_epochs": 1,
        "training#num_epochs": 2,
        "training#num_epochs_per_validation": 1,
        "searching#num_epochs": 2,
        "searching#num_epochs_per_validation": 1,
        "searching#num_warmup_epochs": 1,
    }
    search_space = {"learning_rate": {"_type": "choice", "_value": [0.0001, 0.01]}}
    runner.set_num_fold(num_fold=1)
    runner.set_hpo_params(params=hpo_params)
    runner.set_nni_search_space(search_space)
    runner.run()