# ad hoc: Scaled Logistic Regression with probabilistic output

class force_predict(object):
    def __init__(self, clf, mode='predict_proba', axis=0):
        self._mode = mode
        self._axis = axis
        self._clf = clf

    def fit(self, X, y, **kwargs):
        self._clf.fit(X, y, **kwargs)
        self._copyattr()

    def predict(self, X):
        if self._mode == 'predict_proba':
            return self._clf.predict_proba(X)[:, self._axis]
        elif self._mode == 'decision_function':
            distances = self._clf.decision_function(X)
            if len(distances.shape) > 1:
                return distances[:, self._axis]
            else:
                return distances
        else:
            return self._clf.predict(X)

    def get_params(self, deep=True):
        return dict(clf=self._clf, mode=self._mode, axis=self._axis)

    def _copyattr(self):
        for key, value in self._clf.__dict__.iteritems():
            self.__setattr__(key, value)

from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.linear_model import LogisticRegression

scaler = StandardScaler()
clf = force_predict(LogisticRegression(penalty='l2', C=1), axis=1)
pipeline = Pipeline([('scaler', scaler), ('clf', clf)])


# Area Under the Curve Scorer:
def scorer_auc(y_true, y_pred):
    from sklearn.metrics import roc_auc_score
    from sklearn.preprocessing import LabelBinarizer
    le = LabelBinarizer()
    y_true = le.fit_transform(y_true)
    return roc_auc_score(y_true, y_pred)

# GAT
from mne.decoding import GeneralizationAcrossTime
gat = GeneralizationAcrossTime(n_jobs=-1, clf=pipeline, scorer=auc_scorer)