#
# Copyright 2018 Analytics Zoo Authors.
#
# 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 ray
from ray import tune
import os
from zoo.orca.automl.search.base import SearchEngine, TrialOutput, GoodError
from zoo.orca.automl.search.parameters import DEFAULT_LOGGER_NAME, DEFAULT_METRIC_NAME
from zoo.orca.automl.search.ray_tune.utils import convert_bayes_configs
from zoo.orca.automl.search.utils import get_ckpt_hdfs, put_ckpt_hdfs
from zoo.orca.automl.search import TensorboardLogger
from ray.tune import Stopper
from zoo.orca.automl.model.abstract import ModelBuilder


class RayTuneSearchEngine(SearchEngine):
    """
    Tune driver
    """

    def __init__(self,
                 logs_dir="",
                 resources_per_trial=None,
                 name="",
                 remote_dir=None,
                 ):
        """
        Constructor
        :param logs_dir: local dir to save training results
        :param resources_per_trial: resources for each trial
        :param name: searcher name
        :param remote_dir: checkpoint will be uploaded to remote_dir in hdfs
        """
        self.train_func = None
        self.resources_per_trial = resources_per_trial
        self.trials = None
        self.name = name
        self.remote_dir = remote_dir or RayTuneSearchEngine.get_default_remote_dir(name)
        self.logs_dir = os.path.abspath(os.path.expanduser(logs_dir))

    @staticmethod
    def get_default_remote_dir(name):
        from zoo.ray import RayContext
        from zoo.orca.automl.search.utils import process
        ray_ctx = RayContext.get()
        if ray_ctx.is_local:
            return None
        else:
            default_remote_dir = f"hdfs:///tmp/{name}"
            process(command=f"hadoop fs -mkdir -p {default_remote_dir}")
            return default_remote_dir

    def compile(self,
                data,
                model_builder,
                metric_mode,
                epochs=1,
                validation_data=None,
                metric=None,
                metric_threshold=None,
                n_sampling=1,
                search_space=None,
                search_alg=None,
                search_alg_params=None,
                scheduler=None,
                scheduler_params=None,
                mc=False):
        """
        Do necessary preparations for the engine
        :param data: data for training
               Pandas Dataframe:
                   a Pandas dataframe for training
               Numpy ndarray:
                   a tuple in form of (x, y)
                        x: ndarray for training input
                        y: ndarray for training output
        :param model_builder: model creation function
        :param epochs: max epochs for training
        :param validation_data: validation data
        :param metric: metric name or metric function
        :param metric_mode: mode for metric. "min" or "max". We would infer metric_mode automated
            if user used our built-in metric in zoo.automl.common.metric.Evaluator.
        :param metric_threshold: a trial will be terminated when metric threshold is met
        :param n_sampling: number of sampling
        :param search_space: a dictionary of search_space
        :param search_alg: str, all supported searcher provided by ray tune
               (i.e."variant_generator", "random", "ax", "dragonfly", "skopt",
               "hyperopt", "bayesopt", "bohb", "nevergrad", "optuna", "zoopt" and
               "sigopt")
        :param search_alg_params: extra parameters for searcher algorithm
        :param scheduler: str, all supported scheduler provided by ray tune
        :param scheduler_params: parameters for scheduler
        :param mc: if calculate uncertainty
        """
        # metric and metric's mode
        self.metric_name = metric.__name__ if callable(metric) else (metric or DEFAULT_METRIC_NAME)
        self.mode = metric_mode
        self.stopper = TrialStopper(metric_threshold=metric_threshold,
                                    epochs=epochs,
                                    metric=self.metric_name,
                                    mode=self.mode)
        self.num_samples = n_sampling
        self.search_space = search_space

        self._search_alg = RayTuneSearchEngine._set_search_alg(search_alg, search_alg_params,
                                                               self.metric_name, self.mode)
        self._scheduler = RayTuneSearchEngine._set_scheduler(scheduler, scheduler_params,
                                                             self.metric_name, self.mode)
        metric_func = None if not callable(metric) else metric
        self.train_func = self._prepare_train_func(data=data,
                                                   model_builder=model_builder,
                                                   validation_data=validation_data,
                                                   metric_name=self.metric_name,
                                                   metric_func=metric_func,
                                                   mode=self.mode,
                                                   mc=mc,
                                                   remote_dir=self.remote_dir,
                                                   resources_per_trial=self.resources_per_trial
                                                   )

    @staticmethod
    def _set_search_alg(search_alg, search_alg_params, metric, mode):
        if search_alg:
            if not isinstance(search_alg, str):
                raise ValueError(f"search_alg should be of type str."
                                 f" Got {search_alg.__class__.__name__}")
            params = search_alg_params.copy() if search_alg_params else dict()
            if metric and "metric" not in params:
                params["metric"] = metric
            if mode and "mode" not in params: