"""
Various postprocessing classes implemented as scikit-learn
transformers compatible with scikit-learn pipelines.
"""

import numpy as np

from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.preprocessing import LabelEncoder
from sklearn.utils.validation import check_is_fitted
from sklearn.utils import Bunch

__all__ = [
    "SimpleVoter"
    ]

class SimpleVoter(BaseEstimator, ClassifierMixin):
    """
    Simple voting classifier which takes fitted estimators as
    input and applies voting ensembling at predict time. Similar to
    `sklearn.enseble.VotingClassifier` except the VotingClassifier
    will actually fit the child estimators, where the SimpleVoter
    will take fitted estimators at instantiation and apply similar
    prediction functions as predict time. 
    
    The VotingClassifier implements a dummy fit method that does nothing
    unless the class attributes `estimators` or `classes` have been 
    manually updated.
    
    The value add over `sklearn.enseble.VotingClassifier` is that
    SimpleVoter can solely run the predict, and let the fit live elsewhere.
    
    Args:
        estimators (list of tuples): list of fitted (name, estimator)
            tuples for voting
        classes (array-like): list of class names from one
            of the underlying estimators
        voting (str): voting method ('hard' or 'soft')
        weights (array-like): array of weights corresponding
            to each estimator
    
    """
    def __init__(self, estimators, classes, voting='hard', weights=None):
        self.estimators = estimators
        self.classes = classes
        self.voting = voting
        self.weights = weights
        self._assemble_attributes()
        
    @property
    def named_estimators(self):
        """ Bunches the estimators by name """
        return Bunch(**dict(self.estimators))
        
    @property
    def _weights_not_none(self):
        """Get the weights of not `None` estimators"""
        if self.weights is None:
            return None
        return [w for est, w in zip(self.estimators, self.weights)
                if est[1] not in (None, 'drop')]
        
    def fit(self, X, y=None):
        """ Trivial fit method; re-assembles attributes """
        self._assemble_attributes()
        return self

    def predict(self, X):
        """ Compute predictions for samples in X """
        check_is_fitted(self, 'estimators_')
        if self.voting == 'soft':
            maj = np.argmax(self.predict_proba(X), axis=1)
        else:  # 'hard' voting
            predictions = self._predict(X)
            maj = np.apply_along_axis(
                lambda x: np.argmax(
                    np.bincount(x, weights=self._weights_not_none)),
                axis=1, arr=predictions)
        maj = self.le_.inverse_transform(maj)
        return maj
        
    @property
    def predict_proba(self):
        """ Compute probabilities of possible outcomes for samples in X """
        return self._predict_proba
        
    def _predict(self, X):
        "" "Collect results from clf.predict calls """
        return np.asarray([clf.predict(X) for clf in self.estimators_]).T
        
    def _predict_proba(self, X):
        """ Predict class probabilities for X in 'soft' voting """
        if self.voting == 'hard':
            raise AttributeError("predict_proba is not available when"
                                 " voting=%r" % self.voting)
        check_is_fitted(self, 'estimators_')
        avg = np.average(self._collect_probas(X), axis=0,
                         weights=self._weights_not_none)
        return avg
        
    def _collect_probas(self, X):
        """Collect results from clf.predict calls. """
        return np.asarray([clf.predict_proba(X) for clf in self.estimators_])
        
    def _assemble_attributes(self):
        """ Assemble fitted class attributes """
        names, clfs = zip(*self.estimators)
        self.estimators_ = clfs
        self.classes_ = self.classes
        self.le_ = LabelEncoder()
        self.le_.classes_ = self.classes
