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

__all__ = ['ggdBenchmark', 'ClassificationBenchmark', 'RegressionBenchmark']

# Cell
#export
import pandas as pd
from sklearn.base import BaseEstimator
from sklearn.model_selection import TimeSeriesSplit, StratifiedShuffleSplit, GridSearchCV
from sklearn.pipeline import Pipeline
from sklearn.utils.metaestimators import available_if
from sklearn.utils.validation import check_is_fitted
from .model_documentation import ModelCard

def _benchmark_has(attr):
    """This function is used in `ggdBenchmark` to check if the benchmark has certain attributes"""
    def check(self):
        getattr(self.benchmark, attr)
        return True
    return check

class ggdBenchmark(BaseEstimator):
    """
    The base class for gingado's Benchmark objects.
    """
    def _check_is_time_series(self, X, y=None):
        """
        Checks whether the data is a time series, and sets a data splitter
        accordingly if no data splitter is provided by the user
        Note: all data without an index (eg, a Numpy array) are considered to NOT be a time series
        """
        if hasattr(X, "index"):
            self.is_timeseries = pd.core.dtypes.common.is_datetime_or_timedelta_dtype(X.index)
        else:
            self.is_timeseries = False
        if self.is_timeseries and y is not None:
            if hasattr(y, "index"):
                self.is_timeseries = pd.core.dtypes.common.is_datetime_or_timedelta_dtype(y.index)
            else:
                self.is_timeseries = False

        if self.cv is None:
            self.cv = TimeSeriesSplit() if self.is_timeseries else StratifiedShuffleSplit()

    def _creates_estimator(self):
        if self.estimator is None:
            pass

    def _fit(self, X, y):
        self._check_is_time_series(X, y)

        X, y = self._validate_data(X, y)

        if hasattr(self.estimator, "random_state"):
            self.estimator.random_state = self.random_state

        if self.param_search and self.param_grid:
            self.benchmark = self.param_search(estimator=self.estimator, param_grid=self.param_grid, scoring=self.scoring, verbose=self.verbose_grid)
        else:
            self.benchmark = self.estimator

        self.benchmark.fit(X, y)

        if self.auto_document is not None:
            self.document()

        return self

    def set_benchmark(self, estimator):
        check_is_fitted(estimator)
        self.benchmark = estimator

    def _read_candidate_params(self, candidates, ensemble_method):
                param_grid = []
                for i, model in enumerate(candidates):
                    check_is_fitted(model)
                    param_grid.append({
                        **{'candidate_estimator': [model]},
                        **{
                            'candidate_estimator__' + k: (v,)
                            for k, v in model.get_params().items()
                        }}
                    )
                if ensemble_method is not None:
                    candidate_models = [('candidate_'+str(i+1), model) for i, model in enumerate(candidates)]
                    voting = ensemble_method(estimators=candidate_models)
                    ensemble = {'candidate_estimator': [voting]}
                    param_grid.append(ensemble)
                return param_grid

    def compare(self, X, y, candidates, ensemble_method='object_default', update_benchmark=True):
        """
        Uses a test dataset to compare the performance of the fitted benchmark model with one or more candidate models
        This method achieves this by conducting a grid search
        """
        check_is_fitted(self.benchmark)
        old_benchmark_params = self.benchmark.get_params()

        candidates = list(candidates) if type(candidates) != list else candidates
        list_candidates = [self.benchmark] + candidates

        est = self.benchmark.best_estimator_ if hasattr(self.benchmark, "best_estimator_") else self.benchmark
        cand_pipeline = Pipeline([('candidate_estimator', est)])

        if ensemble_method == 'object_default':
            ensemble_method = self.ensemble_method
        cand_params = self._read_candidate_params(list_candidates, ensemble_method=ensemble_method)
        cand_grid = GridSearchCV(cand_pipeline, cand_params, verbose=self.verbose_grid).fit(X, y)

        self.model_comparison_ = cand_grid

        if update_benchmark:
            if cand_grid.best_estimator_.get_params() != old_benchmark_params:
                self.set_benchmark(cand_grid)
                print("Benchmark updated!")
                print("New benchmark:")
                print(self.benchmark.best_estimator_)

        if self.auto_document is not None:
            self.document()

    def compare_fitted_candidates(self, X, y, candidates, scoring_func):
        check_is_fitted(self.benchmark)
        candidates = list(candidates) if type(candidates) != list else candidates
        for candidate in candidates:
            check_is_fitted(candidate)
        list_candidates = [self.benchmark] + candidates

        return {candidate.__repr__(): scoring_func(y, candidate.predict(X)) for candidate in list_candidates}

    def _read_attr(self):
        for a in dir(self.benchmark):
            if a == '_estimator_type' or a.endswith("_") and not a.startswith("_") and not a.endswith("__"):
                try:
                    model_attr = self.benchmark.__getattribute__(a)
                    yield {a: model_attr}
                except:
                    pass

    def document(self, documenter=None):
        documenter = self.auto_document if documenter is None else documenter
        self.model_documentation = documenter()
        model_info = list(self._read_attr())
        model_info = {k:v for i in model_info for k, v in i.items()}
        #self.model_documentation.read_model(self.benchmark)
        self.model_documentation.fill_model_info(model_info)

    @available_if(_benchmark_has("predict"))
    def predict(self, X, **predict_params):
        return self.benchmark.predict(X, **predict_params)

    @available_if(_benchmark_has("fit_predict"))
    def fit_predict(self, X, y=None, **predict_params):
        return self.benchmark.fit_predict(X, y, **predict_params)

    @available_if(_benchmark_has("predict_proba"))
    def predict_proba(self, X, **predict_proba_params):
        return self.benchmark.predict_proba(X, **predict_proba_params)

    @available_if(_benchmark_has("decision_function"))
    def decision_function(self, X):
        return self.benchmark.decision_function(X)

    @available_if(_benchmark_has("score"))
    def score(self, X):
        return self.benchmark.score(X)

    @available_if(_benchmark_has("score_samples"))
    def score_samples(self, X):
        return self.benchmark.score_samples(X)

    @available_if(_benchmark_has("predict_log_proba"))
    def predict_log_proba(self, X, **predict_log_proba_params):
        return self.benchmark.predict_log_proba(X, **predict_log_proba_params)

# Cell
#export
from .model_documentation import ModelCard
from sklearn.base import ClassifierMixin
from sklearn.model_selection import GridSearchCV
from sklearn.ensemble import RandomForestClassifier, VotingClassifier

class ClassificationBenchmark(ggdBenchmark, ClassifierMixin):
    def __init__(self,
    cv=None,
    estimator=RandomForestClassifier(oob_score=True),
    param_grid={'n_estimators': [100, 250], 'max_features': ['sqrt', 'log2', None]},
    param_search=GridSearchCV,
    scoring=None,
    auto_document=ModelCard,
    random_state=None,
    verbose_grid=False,
    ensemble_method=VotingClassifier):
        self.cv = cv
        self.estimator = estimator
        self.param_grid = param_grid
        self.param_search = param_search
        self.scoring = scoring
        self.auto_document = auto_document
        self.random_state = random_state
        self.verbose_grid = verbose_grid
        self.ensemble_method = ensemble_method

    def fit(self, X, y=None):
        self._fit(X, y)
        return self

# Cell
#export
from sklearn.base import RegressorMixin
from sklearn.ensemble import RandomForestRegressor, VotingRegressor

class RegressionBenchmark(ggdBenchmark, RegressorMixin):
    def __init__(self,
    cv=None,
    estimator=RandomForestRegressor(oob_score=True),
    param_grid={'n_estimators': [100, 250], 'max_features': ['sqrt', 'log2', None]},
    param_search=GridSearchCV,
    scoring=None,
    auto_document=ModelCard,
    random_state=None,
    verbose_grid=False,
    ensemble_method=VotingRegressor):
        self.cv = cv
        self.estimator = estimator
        self.param_grid = param_grid
        self.param_search = param_search
        self.scoring = scoring
        self.auto_document = auto_document
        self.random_state = random_state
        self.verbose_grid = verbose_grid
        self.ensemble_method = ensemble_method

    def fit(self, X, y=None):
        self._fit(X, y)
        return self