# AUTOGENERATED! DO NOT EDIT! File to edit: 01_bayes_opt.ipynb (unless otherwise specified).

__all__ = ['optimize_bayes_param']

# Cell
from sklearn.model_selection import cross_val_score
from bayes_opt import BayesianOptimization
from bayes_opt import JSONLogger
from bayes_opt.event import Events
from bayes_opt.util import load_logs
from pathlib import Path
from functools import wraps
from typing import Dict


# Cell
def optimize_bayes_param(X, y,  *args_objective, objective_fn=cross_val_score, **kwargs_objective):
    """
    The first closure passes in X, y and the objective funtion, which defaults to the cross_val_score function from sklearn.


    :param objective_fn:
        objective function to optimize
    :param X: np.array
        Matrix of features
    :param y: np.array
        Vectors of labels
    :param args_eval:
        *args passed to objective_fn
    :param kwargs_eval:
        **kwargs passed to objective_fn
    """

    def optimize_bayes_wo_param(parse_model_params):

        def _opt_engine(*args_model, **kwargs_model):
            """
            This is the running engine function that takes an estimator/model object and a loss function

            """

            estimator = parse_model_params(*args_model, **kwargs_model)
            return objective_fn(estimator, X=X, y=y, *args_objective, **kwargs_objective).mean()

        @wraps(parse_model_params)
        def run_trials(pbounds: Dict,
                       init_points: int = 10,
                       n_iter: int = 10,
                       log_dir: Path = Path("./bayes_opt_logs"),
                       acq: str = 'ucb',
                       kappa: str = 2.576,
                       fit: bool = True):
            """
            :param pbounds: dict
                Dictionary with parameters names as keys and a tuple with minimum
                and maximum values.
            :param init_points : int
                Number of iterations before the explorations starts the exploration
                for the maximum.
            :param n_iter: int
                Number of iterations where the method attempts to find the maximum
                value.
            :param log_dir: Path
                Directory to log json results
            :param acq: str
                The acquisition method used.
            :param kappa: float
                Parameter to indicate how closed are the next parameters sampled.
                Higher value = favors spaces that are least explored.
                Lower value = favors spaces where the regression function is the
                highest.
            :param fit: bool
                if True the best model is fitted on de data
                if False the best model is returned unfitted
            :return: A Sklearn model with the pbounds hyperparameters optimized
             by Bayesian Optimisation in cross-validation
            """
            optimizer = BayesianOptimization(_opt_engine, pbounds=pbounds)
            log_dir = Path(log_dir)
            if log_dir.exists():
                all_log = [str(path) for path in log_dir.iterdir()]
                load_logs(optimizer, logs=all_log)
                filename = 'log_{}.json'.format(len(all_log))
            else:
                log_dir.mkdir(parents=True)
                filename = 'log_0.json'
            logger = JSONLogger(path=str(log_dir / filename))
            optimizer.subscribe(Events.OPTIMIZATION_STEP, logger)

            optimizer.maximize(init_points, n_iter, kappa=kappa, acq=acq)
            print(f"The best combination of hyperparameters are { optimizer.max['params'] }")
            best_model = parse_model_params(**optimizer.max['params'])
            if fit:
                best_model.fit(X=X, y=y)
            return best_model

        return run_trials

    return optimize_bayes_wo_param