import warnings

import numpy as np
from sklearn.base import BaseEstimator
from sklearn.externals.joblib import delayed
from sklearn.externals.joblib import Parallel


warnings.warn("""
The nolearn.model module will be removed in nolearn 0.6.  If you want
to continue to use this module, please consider copying the code into
your own project.
""")


class AbstractModel(object):
    """A small abstraction around :class:`~sklearn.pipeline.Pipeline`
    objects.

    Allows the convenient parametrization of the underlying pipeline
    through :attr:`~AbstractModel.params`.
    """
    default_params = dict()

    def __init__(self, **kwargs):
        """
        :param kwargs: Keyword arguments correspond to pipeline
                       parameters, and will override parameters in
                       :attr:`~AbstractModel.default_params`.
        """
        params = self.default_params.copy()
        params.update(kwargs)
        self.params = params

    def __call__(self):
        """
        :rtype: :class:`~sklearn.pipeline.Pipeline`
        """
        pipeline = self.pipeline
        pipeline.set_params(**self.params)
        return pipeline

    @property
    def pipeline(self):  # pragma: no cover
        raise NotImplementedError()


def _avgest_fit_est(est, i, X, y, verbose):
    if verbose:
        print "[AveragingEstimator] estimator_%s.fit() ..." % i
    return est.fit(X, y)


def _avgest_predict_proba(est, i, X, verbose):
    if verbose:
        print "[AveragingEstimator] estimator_%s.predict_proba() ..." % i
    return est.predict_proba(X)


class AveragingEstimator(BaseEstimator):
    """An estimator that wraps a list of other estimators and returns
    their average for :meth:`fit`, :meth:`predict` and
    :meth:`predict_proba`.
    """
    def __init__(self, estimators, verbose=0, n_jobs=1):
        """
        :param estimators: List of estimator objects.
        """
        self.estimators = estimators
        self.verbose = verbose
        self.n_jobs = n_jobs

    def fit(self, X, y):
        result = Parallel(n_jobs=self.n_jobs, verbose=self.verbose)(
            delayed(_avgest_fit_est)(est, i, X, y, self.verbose)
            for i, est in enumerate(self.estimators))
        self.estimators = result
        return self

    def predict(self, X):
        return np.argmax(self.predict_proba(X), axis=1)

    def predict_proba(self, X):
        result = Parallel(n_jobs=self.n_jobs, verbose=self.verbose)(
            delayed(_avgest_predict_proba)(est, i, X, self.verbose)
            for i, est in enumerate(self.estimators))
        for proba in result[1:]:
            result[0] += proba
        return result[0] / len(self.estimators)
