# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/07_sklearn.ipynb (unless otherwise specified).

__all__ = ['ScikitLearner', 'ScikitLearnerCV']

# Cell
import sklearn

from sklearn.metrics import *
from functools import lru_cache, partial


class ScikitLearner:
    """Helper class to use fine tune parameter based on validation dataset performace.
    Feel free to implement your own version."""
    def __init__(self, learner, X_train, X_val, y_train, y_val, predict_proba=True):
        """Create the class instance.

        **Parameters**

        - learner: the scikit-learn estimator
        - X_train, X_val, y_train, y_val: train validation datasets
        - predict_proba: whether the estimator can predict probability
        """
        self.learner = learner
        self.predict_proba = predict_proba
        self._proba = None
        self.X_train, self.X_val, self.y_train, self.y_val = X_train, X_val, y_train, y_val

    def __getattr__(self, key):
        return getattr(self.learner, key)

    def __dir__(self):
        return set(super().__dir__() + list(self.__dict__.keys()) +
                   dir(self.learner))

    @lru_cache(maxsize=None)
    def valid_loss_with_params(self, loss_func=None, callbacks=tuple(), **params):
        """Calculate loss of the estimator on validation set, and also the values of callbacks.

        **Parameters**

        - loss_func: loss function to use. By default, if estimator support to predict probability, it will be
        log_loss, otherwise it will be accuracy score.
        - callbacks: callbacks to also evaluate. Default is empty tuple.
        - params: parameters to use when train the estimator.

        **Returns**

        Return values will be a two element tuple.
        - the first is values of the loss function
        - the other is values of all callbacks
        """
        return self.loss_with_params(self.X_val,
                                     self.y_val,
                                     loss_func=loss_func,
                                     callbacks=callbacks,
                                     **params)

    def loss_with_params(self, X, y, loss_func=None, callbacks=tuple(), **params):
        estimator = sklearn.clone(self.learner)
        estimator.set_params(**params)
        estimator.fit(self.X_train, self.y_train)
        if self.predict_proba:
            if self._proba is None:
                try:
                    y_pred = estimator.predict_proba(self.X_val)
                    self._proba = True
                except:
                    self._proba = False
                    y_pred = estimator.predict(self.X_val)
            elif self._proba:
                y_pred = estimator.predict_proba(self.X_val)
            else:
                y_pred = estimator.predict(self.X_val)
        else:
            y_pred = estimator.predict(self.X_val)
        if loss_func is None:
            if self._proba:
                loss_func = log_loss
            else:
                loss_func = lambda x,y: -accuracy_score(x,y)
        return [loss_func(self.y_val, y_pred)], [callback(self.y_val, y_pred) for callback in callbacks]


# Cell
import sklearn

from sklearn.metrics import *
from functools import lru_cache, partial
from sklearn.model_selection import cross_val_predict


class ScikitLearnerCV:
    "Make use of sklearn cross_val_predict interface to optimize paramters."

    def __init__(self, learner, X, y):
        """Create the class instance.

        **Parameters**

        - learner: the scikit-learn estimator
        """
        self.learner = learner
        self.X = X
        self.y = y

    def __getattr__(self, key):
        return getattr(self.learner, key)

    def __dir__(self):
        return set(super().__dir__() + list(self.__dict__.keys()) +
                   dir(self.learner))

    @lru_cache(maxsize=None)
    def cv_loss_with_params(self,
                            loss_func=None,
                            callbacks=tuple(),
                            groups=None,
                            cv=None,
                            n_jobs=None,
                            verbose=0,
                            pre_dispatch='2*n_jobs',
                            method='predict',
                            fit_params=None,
                            **params):
        """Calculate loss of the estimator on validation set, and also the values of callbacks.

        **Parameters**

        - loss_func: loss function to use. By default, if estimator support to predict probability, it will be
        log_loss, otherwise it will be accuracy score.
        - callbacks: callbacks to also evaluate. Default is empty tuple.
        - params: parameters to use when train the estimator.

        **Returns**

        Return values will be a two element tuple.
        - the first is values of the loss function
        - the other is values of all callbacks
        """
        return self.loss_with_params(self.X,
                                     self.y,
                                     loss_func=loss_func,
                                     callbacks=callbacks,
                                     groups=groups,
                                     cv=cv,
                                     n_jobs=n_jobs,
                                     verbose=verbose,
                                     pre_dispatch=pre_dispatch,
                                     method=method,
                                     fit_params=fit_params,
                                     **params)

    def loss_with_params(self,
                         X,
                         y,
                         loss_func=None,
                         callbacks=tuple(),
                         groups=None,
                         cv=None,
                         n_jobs=None,
                         verbose=0,
                         pre_dispatch='2*n_jobs',
                         method='predict',
                         fit_params=None,
                         **params):
        estimator = sklearn.clone(self.learner)
        estimator.set_params(**params)
        y_pred = cross_val_predict(estimator,
                                   X,
                                   y,
                                   groups=groups,
                                   cv=cv,
                                   n_jobs=n_jobs,
                                   verbose=verbose,
                                   pre_dispatch=pre_dispatch,
                                   fit_params=fit_params,
                                   method=method)
        if loss_func is None:
            if method == 'predict_proba':
                loss_func = log_loss
            else:
                loss_func = lambda x, y: -accuracy_score(x, y)
        return [loss_func(y, y_pred)
                ], [callback(y, y_pred) for callback in callbacks]