# Copyright (c) 2019 Uber Technologies, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import contextlib
import logging
import os.path

import mlflow
import pytest
import ray
from mlflow.tracking import MlflowClient

from ludwig.constants import ACCURACY, TRAINER
from ludwig.contribs import MlflowCallback
from ludwig.hyperopt.execution import get_build_hyperopt_executor
from ludwig.hyperopt.results import RayTuneResults
from ludwig.hyperopt.run import hyperopt, update_hyperopt_params_with_defaults
from ludwig.hyperopt.sampling import get_build_hyperopt_sampler
from ludwig.utils.defaults import merge_with_defaults
from tests.integration_tests.utils import category_feature, generate_data, spawn, text_feature

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logging.getLogger("ludwig").setLevel(logging.INFO)

HYPEROPT_CONFIG = {
    "parameters": {
        "trainer.learning_rate": {
            "space": "loguniform",
            "lower": 0.001,
            "upper": 0.1,
        },
        "combiner.num_fc_layers": {"space": "randint", "lower": 2, "upper": 6},
        "utterance.cell_type": {"space": "grid_search", "values": ["rnn", "gru"]},
        "utterance.bidirectional": {"space": "choice", "categories": [True, False]},
        "utterance.fc_layers": {
            "space": "choice",
            "categories": [
                [{"output_size": 64}, {"output_size": 32}],
                [{"output_size": 64}],
                [{"output_size": 32}],
            ],
        },
    },
    "goal": "minimize",
}


SAMPLERS = [
    {"type": "ray"},
    {"type": "ray", "num_samples": 2},
    {
        "type": "ray",
        "search_alg": {"type": "bohb"},
        "scheduler": {
            "type": "hb_bohb",
            "time_attr": "training_iteration",
            "reduction_factor": 4,
        },
        "num_samples": 3,
    },
]

EXECUTORS = [
    {"type": "ray"},
]


def _get_config(sampler, executor):
    input_features = [
        text_feature(name="utterance", cell_type="lstm", reduce_output="sum"),
        category_feature(vocab_size=2, reduce_input="sum"),
    ]

    output_features = [category_feature(vocab_size=2, reduce_input="sum")]

    return {
        "input_features": input_features,
        "output_features": output_features,
        "combiner": {"type": "concat", "num_fc_layers": 2},
        TRAINER: {"epochs": 2, "learning_rate": 0.001},
        "hyperopt": {
            **HYPEROPT_CONFIG,
            "executor": executor,
            "sampler": sampler,
        },
    }


@contextlib.contextmanager
def ray_start_4_cpus():
    res = ray.init(
        num_cpus=4,
        include_dashboard=False,
        object_store_memory=150 * 1024 * 1024,
    )
    try:
        yield res
    finally:
        ray.shutdown()


@spawn
def run_hyperopt_executor(
    sampler,
    executor,
    csv_filename,
    validate_output_feature=False,
    validation_metric=None,
):
    config = _get_config(sampler, executor)
    rel_path = generate_data(config["input_features"], config["output_features"], csv_filename)

    config = merge_with_defaults(config)

    hyperopt_config = config["hyperopt"]

    if validate_output_feature:
        hyperopt_config["output_feature"] = config["output_features"][0]["name"]
    if validation_metric:
        hyperopt_config["validation_metric"] = validation_metric

    update_hyperopt_params_with_defaults(hyperopt_config)

    parameters = hyperopt_config["parameters"]
    if sampler.get("search_alg", {}).get("type", "") == "bohb":
        # bohb does not support grid_search search space
        del parameters["utterance.cell_type"]

    split = hyperopt_config["split"]
    output_feature = hyperopt_config["output_feature"]
    metric = hyperopt_config["metric"]
    goal = hyperopt_config["goal"]

    hyperopt_sampler = get_build_hyperopt_sampler(sampler["type"])(goal, parameters, **sampler)

    hyperopt_executor = get_build_hyperopt_executor(executor["type"])(
        hyperopt_sampler, output_feature, metric, split, **executor
    )

    hyperopt_executor.execute(
        config,
        dataset=rel_path,
        backend="local",
    )


@pytest.mark.distributed
@pytest.mark.parametrize("sampler", SAMPLERS)
@pytest.mark.parametrize("executor", EXECUTORS)
def test_hyperopt_executor(sampler, executor, csv_filename):
    with ray_start_4_cpus():
        run_hyperopt_executor(sampler, executor, csv_filename)


@pytest.mark.distributed
def test_hyperopt_executor_with_metric(csv_filename):
    with ray_start_4_cpus():
        run_hyperopt_executor(
            {"type": "ray", "num_samples": 2},
            {"type": "ray"},
            csv_filename,
            validate_output_feature=True,
            validation_metric=ACCURACY,
        )


@pytest.mark.distributed
def test_hyperopt_run_hyperopt(csv_filename):
    with ray_start_4_cpus():
        input_features = [
            text_feature(name="utterance", cell_type="lstm", reduce_output="sum"),
            category_feature(vocab_size=2, reduce_input="sum"),
        ]

        output_features = [category_feature(vocab_size=2, reduce_input="sum")]

        rel_path = generate_data(input_features, output_features, csv_filename)

        config = {
            "input_features": input_features,
            "output_features": output_features,
            "combiner": {"type": "concat", "num_fc_layers": 2},
            TRAINER: {"epochs": 2, "learning_rate": 0.001},
        }

        output_feature_name = output_features[0]["name"]

        hyperopt_configs = {
            "parameters": {
                "trainer.learning_rate": {
                    "space": "loguniform",
                    "lower": 0.001,
                    "upper": 0.1,
                },
                output_feature_name + ".output_size": {"space": "randint", "lower": 32, "upper": 64},
                output_feature_name + ".num_fc_layers": {"space": "randint", "lower": 2, "upper": 6},
            },
            "goal": "minimize",
            "output_feature": output_feature_name,
            "validation_metrics": "loss",
            "executor": {"type": "ray"},
            "sampler": {"type": "ray", "num_samples": 2},
        }

        # add hyperopt parameter space to the config
        config["hyperopt"] = hyperopt_configs
        run_hyperopt(config, rel_path)


@pytest.mark.distributed
def test_hyperopt_ray_mlflow(csv_filename, tmpdir):
    with ray_start_4_cpus():
        mlflow_uri = f"file://{tmpdir}/mlruns"
        mlflow.set_tracking_uri(mlflow_uri)
        client = MlflowClient(tracking_uri=mlflow_uri)

        num_samples = 2
        config = _get_config({"type": "ray", "num_samples": num_samples}, {"type": "ray"})

        rel_path = generate_data(config["input_features"], config["output_features"], csv_filename)

        exp_name = "mlflow_test"
        run_hyperopt(config, rel_path, experiment_name=exp_name, callbacks=[MlflowCallback(mlflow_uri)])

        experiment = client.get_experiment_by_name(exp_name)
        assert experiment is not None

        runs = client.search_runs([experiment.experiment_id])
        assert len(runs) > 0

        for run in runs:
            artifacts = [f.path for f in client.list_artifacts(run.info.run_id, "")]
            assert "config.yaml" in artifacts
            assert "model" in artifacts


@spawn
def run_hyperopt(
    config,
    rel_path,
    experiment_name="ray_hyperopt",
    callbacks=None,
):
    hyperopt_results = hyperopt(
        config,
        dataset=rel_path,
        output_directory="results_hyperopt",
        experiment_name=experiment_name,
        callbacks=callbacks,
    )

    # check for return results
    assert isinstance(hyperopt_results, RayTuneResults)

    # check for existence of the hyperopt statistics file
    assert os.path.isfile(os.path.join("results_hyperopt", "hyperopt_statistics.json"))
