# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

import json
import os
import requests
import subprocess
import sys
import time


NNI_REST_ENDPOINT = "http://localhost:8080/api/v1/nni"
NNI_STATUS_URL = NNI_REST_ENDPOINT + "/check-status"
NNI_TRIAL_JOBS_URL = NNI_REST_ENDPOINT + "/trial-jobs"
WAITING_TIME = 20
MAX_RETRIES = 10


def get_experiment_status(status_url=NNI_STATUS_URL):
    """Helper method. Gets the experiment status from the REST endpoint. 

    Args:
        status_url (str): URL for the REST endpoint

    Returns:
        dict: status of the experiment
    """
    return requests.get(status_url).json()


def check_experiment_status(wait=WAITING_TIME, max_retries=MAX_RETRIES):
    """Checks the status of the current experiment on the NNI REST endpoint.

    Waits until the tuning has completed.

    Args:
        wait (numeric) : time to wait in seconds
        max_retries (int): max number of retries
    """
    i = 0
    while i < max_retries:
        nni_status = get_experiment_status(NNI_STATUS_URL)
        if nni_status["status"] in ["DONE", "TUNER_NO_MORE_TRIAL"]:
            break
        elif nni_status["status"] not in ["RUNNING", "NO_MORE_TRIAL"]:
            raise RuntimeError(
                "NNI experiment failed to complete with status {} - {}".format(
                    nni_status["status"], nni_status["errors"][0]
                )
            )
        time.sleep(wait)
        i += 1
    if i == max_retries:
        raise TimeoutError("check_experiment_status() timed out")


def check_stopped(wait=WAITING_TIME, max_retries=MAX_RETRIES):
    """Checks that there is no NNI experiment active (the URL is not accessible). 
    This method should be called after `nnictl stop` for verification.

    Args:
        wait (numeric) : time to wait in seconds
        max_retries (int): max number of retries
    """
    i = 0
    while i < max_retries:
        try:
            get_experiment_status(NNI_STATUS_URL)
        except:
            break
        time.sleep(wait)
        i += 1
    if i == max_retries:
        raise TimeoutError("check_stopped() timed out")


def check_metrics_written(wait=WAITING_TIME, max_retries=MAX_RETRIES):
    """Waits until the metrics have been written to the trial logs.
    
    Args:
        wait (numeric) : time to wait in seconds
        max_retries (int): max number of retries
    """
    i = 0
    while i < max_retries:
        all_trials = requests.get(NNI_TRIAL_JOBS_URL).json()
        if all(["finalMetricData" in trial for trial in all_trials]):
            break
        time.sleep(wait)
        i += 1
    if i == max_retries:
        raise TimeoutError("check_metrics_written() timed out")


def get_trials(optimize_mode):
    """Obtain information about the trials of the current experiment via the REST endpoint.

    Args:
        optimize_mode (str): One of "minimize", "maximize". Determines how to obtain the best default metric.

    Returns:
         list: Trials info, list of (metrics, log path)
         dict: Metrics for the best choice of hyperparameters
         dict: Best hyperparameters
         str: Log path for the best trial
    """
    if optimize_mode not in ["minimize", "maximize"]:
        raise ValueError("optimize_mode should equal either minimize or maximize")
    all_trials = requests.get(NNI_TRIAL_JOBS_URL).json()
    trials = [
        (eval(trial["finalMetricData"][0]["data"]), trial["logPath"].split(":")[-1])
        for trial in all_trials
    ]
    sorted_trials = sorted(
        trials, key=lambda x: x[0]["default"], reverse=(optimize_mode == "maximize")
    )
    best_trial_path = sorted_trials[0][1]

    # Read the metrics from the trial directory in order to get the name of the default metric
    with open(os.path.join(best_trial_path, "metrics.json"), "r") as fp:
        best_metrics = json.load(fp)
    with open(os.path.join(best_trial_path, "parameter.cfg"), "r") as fp:
        best_params = json.load(fp)
    return trials, best_metrics, best_params, best_trial_path


def stop_nni():
    """Stop nni experiment"""
    proc = subprocess.run([sys.prefix + "/bin/nnictl", "stop"])
    if proc.returncode != 0:
        raise RuntimeError("'nnictl stop' failed with code %d" % proc.returncode)
    check_stopped()


def start_nni(config_path, wait=WAITING_TIME, max_retries=MAX_RETRIES):
    """Start nni experiment given a configuration yaml file.

    Args:
        config_path (str): Configuration yaml file.
        wait (numeric) : time to wait in seconds
        max_retries (int): max number of retries
    """
    nni_env = os.environ.copy()
    nni_env["PATH"] = sys.prefix + "/bin:" + nni_env["PATH"]
    proc = subprocess.run(["nnictl", "create", "--config", config_path], env=nni_env)
    if proc.returncode != 0:
        raise RuntimeError("'nnictl create' failed with code %d" % proc.returncode)
    check_experiment_status(wait=wait, max_retries=max_retries)
