import copy
import functools
import math
import operator
import typing

from creme import base
from creme import metrics


__all__ = ['SuccessiveHalvingClassifier', 'SuccessiveHalvingRegressor']


class SuccessiveHalving:

    def __init__(self, models: typing.List[base.Predictor], metric: metrics.Metric, budget: int,
                 eta=2, verbose=False, **print_kwargs):

        # Check that the model and the metric are in accordance
        for model in models:
            if not metric.works_with(model):
                raise ValueError(f"{metric.__class__.__name__} metric can't be used to evaluate a " +
                                f'{model.__class__.__name__}')

        self.models = models
        self.metric = metric
        self.budget = budget
        self.eta = eta
        self.verbose = verbose
        self.print_kwargs = print_kwargs

        self._n = len(models)
        self._metrics = [copy.deepcopy(metric) for _ in range(self._n)]
        self._s = self._n
        self._n_rungs = 0
        self._rankings = list(range(self._n))
        self._r = math.floor(budget / (self._s * math.ceil(math.log(self._n, eta))))
        self._budget_used = 0
        self._n_iterations = 0
        self._best_model_idx = 0

        if isinstance(model, base.Classifier) and not metric.requires_labels:
            self._pred_func = lambda model: model.predict_proba_one
        else:
            self._pred_func = lambda model: model.predict_one

    @property
    def best_model(self):
        """The current best model."""
        return self.models[self._best_model_idx]

    def fit_one(self, x, y):

        for i in self._rankings[:self._s]:
            model = self.models[i]
            metric = self._metrics[i]
            y_pred = self._pred_func(model)(x)
            metric.update(y_true=y, y_pred=y_pred)
            model.fit_one(x, y)

            # Check for a new best model
            if i != self._best_model_idx:
                op = operator.gt if self.metric.bigger_is_better else operator.lt
                if op(self._metrics[i].get(), self._metrics[self._best_model_idx].get()):
                    self._best_model_idx = i

        self._n_iterations += 1

        if self._s > 1 and self._n_iterations == self._r:

            self._n_rungs += 1
            self._budget_used += self._s * self._r

            # Update the rankings of the current models based on their respective metric values
            self._rankings[:self._s] = sorted(
                self._rankings[:self._s],
                key=lambda i: self._metrics[i].get(),
                reverse=self.metric.bigger_is_better
            )

            # Determine how many models to keep for the current rung
            cutoff = math.ceil(self._s / self.eta)

            if self.verbose:
                print(
                    '\t'.join((
                        f'[{self._n_rungs}]',
                        f'{self._s - cutoff} removed',
                        f'{cutoff} left',
                        f'{self._r} iterations',
                        f'budget used: {self._budget_used}',
                        f'budget left: {self.budget - self._budget_used}',
                        f'best {self._metrics[self._rankings[0]]}',
                    )),
                    **self.print_kwargs
                )

            # Determine where the next rung is located
            self._s = cutoff
            self._r = math.floor(self.budget / (self._s * math.ceil(math.log(self._n, self.eta))))

        return self


class SuccessiveHalvingRegressor(SuccessiveHalving, base.Regressor):
    """Successive halving algorithm for regression.

    Successive halving is a method for performing model selection without having to train each
    model on all the dataset. At certain points in time (called "rungs"), the worst performing will
    be discarded and the best ones will keep competing between each other. The rung values are
    designed so that at most `budget` model updates will be performed in total.

    If you have `k` combinations of hyperparameters and that your dataset contains `n`
    observations, then the maximal budget you can allocate is:

    $$\\frac{2kn}{eta}$$

    It is recommended that you check this beforehand. This bound can't be checked by the function
    because the size of the dataset is not known. In fact it is potentially infinite, in which case
    the algorithm will terminate once all the budget has been spent.

    If you have a budget of `B`, and that your dataset contains `n` observations, then the
    number of hyperparameter combinations that will spend all the budget and go through all the
    data is:

    $$\\ceil(\\floor(\\frac{B}{2n}) \\times eta)$$

    Parameters:
        models: The models to compare.
        X_y: A stream of (features, target) tuples.
        metric: Metric used for comparing models with.
        budget: Total number of model updates you wish to allocate.
        eta: Rate of elimination. At every rung, `math.ceil(k / eta)` models are kept, where
            `k` is the number of models that have reached the rung. A higher `eta` value will
            focus on less models but will allocate more iterations to the best models.
        verbose: Whether to display progress or not.
        print_kwargs: Extra keyword arguments are passed to the `print` function. For instance,
            this allows providing a `file` argument, which indicates where to output progress.

    Example:

        >>> from creme import datasets
        >>> from creme import linear_model
        >>> from creme import metrics
        >>> from creme import preprocessing
        >>> from creme import model_selection
        >>> from creme import optim

        As an example, let's use successive halving to tune the learning rate of a linear
        regression. We'll first define the model.

        >>> from creme import linear_model
        >>> from creme import preprocessing

        >>> model = (
        ...     preprocessing.StandardScaler() |
        ...     linear_model.LinearRegression(intercept_lr=.1)
        ... )

        Let's now define a grid of parameters which we would like to compare. We'll try
        different optimizers with various learning rates.

        >>> param_grid = model_selection.expand_param_grid({
        ...     'LinearRegression': {
        ...         'optimizer': [
        ...             (optim.SGD, {'lr': [.1, .01, .005]}),
        ...             (optim.Adam, {'beta_1': [.01, .001], 'lr': [.1, .01, .001]}),
        ...             (optim.Adam, {'beta_1': [.1], 'lr': [.001]}),
        ...         ]
        ...     }
        ... })

        Let's now instantiate multiple copies of the model with each parameter configuration.

        >>> models = [model._set_params(params) for params in param_grid]

        We can check how many models this created.

        >>> len(models)
        10

        We can now pass these models to a `SuccessiveHalvingRegressor`. We also need to pick a
        metric to compare the models, and a budget which indicates how many iterations to run
        before picking the best model and discarding the rest.

        >>> sh = model_selection.SuccessiveHalvingRegressor(
        ...     models=models,
        ...     metric=metrics.MAE(),
        ...     budget=2000,
        ...     eta=2,
        ...     verbose=True
        ... )

        A `SuccessiveHalvingRegressor` is also a regressor with a `fit_one` and a `predict_one`
        method, therefore we can use `creme.model_selection.progressive_val_score` to evaluate it.

        >>> model_selection.progressive_val_score(
        ...     X_y=datasets.TrumpApproval(),
        ...     model=sh,
        ...     metric=metrics.MAE()
        ... )
        [1]	5 removed	5 left	50 iterations	budget used: 500	budget left: 1500	best MAE: 4.540491
        [2]	2 removed	3 left	100 iterations	budget used: 1000	budget left: 1000	best MAE: 2.458765
        [3]	1 removed	2 left	166 iterations	budget used: 1498	budget left: 502	best MAE: 1.583751
        [4]	1 removed	1 left	250 iterations	budget used: 1998	budget left: 2	best MAE: 1.147296
        MAE: 0.488387

        We can now view the best model.

        >>> sh.best_model
        Pipeline (
          StandardScaler (),
          LinearRegression (
            optimizer=Adam (
              lr=Constant (
                learning_rate=0.1
              )
              beta_1=0.01
              beta_2=0.999
              eps=1e-08
            )
            loss=Squared ()
            l2=0.
            intercept=39.93843
            intercept_lr=Constant (
              learning_rate=0.1
            )
            clip_gradient=1e+12
            initializer=Zeros ()
          )
        )

    References:
        1. [Jamieson, K. and Talwalkar, A., 2016, May. Non-stochastic best arm identification and hyperparameter optimization. In Artificial Intelligence and Statistics (pp. 240-248).](http://proceedings.mlr.press/v51/jamieson16.pdf)
        2. [Li, L., Jamieson, K., Rostamizadeh, A., Gonina, E., Hardt, M., Recht, B. and Talwalkar, A., 2018. Massively parallel hyperparameter tuning. arXiv preprint arXiv:1810.05934.](https://arxiv.org/pdf/1810.05934.pdf)
        3. [Li, L., Jamieson, K., DeSalvo, G., Rostamizadeh, A. and Talwalkar, A., 2017. Hyperband: A novel bandit-based approach to hyperparameter optimization. The Journal of Machine Learning Research, 18(1), pp.6765-6816.](https://arxiv.org/pdf/1603.06560.pdf)

    """

    def predict_one(self, x):
        return self.best_model.predict_one(x)


class SuccessiveHalvingClassifier(SuccessiveHalving, base.Classifier):
    """Successive halving algorithm for classification.

    Successive halving is a method for performing model selection without having to train each
    model on all the dataset. At certain points in time (called "rungs"), the worst performing will
    be discarded and the best ones will keep competing between each other. The rung values are
    designed so that at most `budget` model updates will be performed in total.

    If you have `k` combinations of hyperparameters and that your dataset contains `n`
    observations, then the maximal budget you can allocate is:

    .. math:: \\frac{2kn}{eta}

    It is recommended that you check this beforehand. This bound can't be checked by the function
    because the size of the dataset is not known. In fact it is potentially infinite, in which case
    the algorithm will terminate once all the budget has been spent.

    If you have a budget of `B`, and that your dataset contains `n` observations, then the
    number of hyperparameter combinations that will spend all the budget and go through all the
    data is:

    .. math:: \\ceil(\\floor(\\frac{B}{(2n)}) \times eta)

    Parameters:
        models: The models to compare.
        X_y: A stream of (features, target) tuples.
        metric: Metric used for comparing models with.
        budget: Total number of model updates you wish to allocate.
        eta: Rate of elimination. At every rung, `math.ceil(k / eta)` models are kept, where
            `k` is the number of models that have reached the rung. A higher `eta` value will
            focus on less models but will allocate more iterations to the best models.
        verbose: Whether to display progress or not.
        print_kwargs: Extra keyword arguments are passed to the `print` function. For instance,
            this allows providing a `file` argument, which indicates where to output progress.

    Examples:

        >>> from creme import datasets
        >>> from creme import linear_model
        >>> from creme import metrics
        >>> from creme import model_selection
        >>> from creme import optim
        >>> from creme import preprocessing

        >>> model = (
        ...     preprocessing.StandardScaler() |
        ...     linear_model.LogisticRegression()
        ... )

        >>> param_grid = model_selection.expand_param_grid({
        ...     'LogisticRegression': {
        ...         'optimizer': [
        ...             (optim.SGD, {'lr': [.1, .01, .005]}),
        ...             (optim.Adam, {'beta_1': [.01, .001], 'lr': [.1, .01, .001]}),
        ...             (optim.Adam, {'beta_1': [.1], 'lr': [.001]}),
        ...         ]
        ...     }
        ... })

        >>> models = [model._set_params(params) for params in param_grid]

        >>> sh = model_selection.SuccessiveHalvingClassifier(
        ...     models=models,
        ...     metric=metrics.Accuracy(),
        ...     budget=2000,
        ...     eta=2,
        ...     verbose=True
        ... )

        >>> model_selection.progressive_val_score(
        ...     X_y=datasets.Phishing(),
        ...     model=sh,
        ...     metric=metrics.ROCAUC()
        ... )
        [1]	5 removed	5 left	50 iterations	budget used: 500	budget left: 1500	best Accuracy: 80.00%
        [2]	2 removed	3 left	100 iterations	budget used: 1000	budget left: 1000	best Accuracy: 84.00%
        [3]	1 removed	2 left	166 iterations	budget used: 1498	budget left: 502	best Accuracy: 86.14%
        [4]	1 removed	1 left	250 iterations	budget used: 1998	budget left: 2	best Accuracy: 84.80%
        ROCAUC: 0.953198

        >>> sh.best_model
        Pipeline (
          StandardScaler (),
          LogisticRegression (
            optimizer=Adam (
              lr=Constant (
                learning_rate=0.01
              )
              beta_1=0.01
              beta_2=0.999
              eps=1e-08
            )
            loss=Log (
              weight_pos=1.
              weight_neg=1.
            )
            l2=0.
            intercept=-0.399002
            intercept_lr=Constant (
              learning_rate=0.01
            )
            clip_gradient=1e+12
            initializer=Zeros ()
          )
        )

    References:
        1. [Jamieson, K. and Talwalkar, A., 2016, May. Non-stochastic best arm identification and hyperparameter optimization. In Artificial Intelligence and Statistics (pp. 240-248).](http://proceedings.mlr.press/v51/jamieson16.pdf)
        2. [Li, L., Jamieson, K., Rostamizadeh, A., Gonina, E., Hardt, M., Recht, B. and Talwalkar, A., 2018. Massively parallel hyperparameter tuning. arXiv preprint arXiv:1810.05934.](https://arxiv.org/pdf/1810.05934.pdf)
        3. [Li, L., Jamieson, K., DeSalvo, G., Rostamizadeh, A. and Talwalkar, A., 2017. Hyperband: A novel bandit-based approach to hyperparameter optimization. The Journal of Machine Learning Research, 18(1), pp.6765-6816.](https://arxiv.org/pdf/1603.06560.pdf)

    """

    def predict_proba_one(self, x):
        return self.best_model.predict_proba_one(x)
