# -*- coding: utf-8 -*-

"""

Created on Sat Aug 15 14:12:12 2015.

@author: rc, alex

"""

import numpy as np

from collections import OrderedDict

from sklearn.base import BaseEstimator, ClassifierMixin

from sklearn.metrics import roc_auc_score

from hyperopt import fmin, tpe, hp

from progressbar import Bar, ETA, Percentage, ProgressBar, RotatingMarker

class WeightedMeanClassifier(BaseEstimator, ClassifierMixin):

    """Weigted mean classifier with AUC optimization."""

    def __init__(self, ensemble, step=0.025, max_evals=100, mean='arithmetic', 

                 verbose=True):

        """Init."""

        self.ensemble = ensemble

        self.step = step

        self.max_evals = max_evals

        self.mean = mean

        self.count = -1

        self.verbose = verbose

        self.param_space = OrderedDict()

        for model in ensemble:

            self.param_space[model] = hp.quniform(model, 0, 3, self.step)

        # input data are arranged in a particular order, whereas hyperopt uses 

        # unordered lists when optimizing. The model has to keep track 

        # of the initial order so that correct weights are applied to columns

        self.sorting = dict()

        for i, m in enumerate(self.ensemble):

            self.sorting[m] = i

    def fit(self, X, y):

        """Fit."""

        self.best_params = None

        if self.mean != 'simple':

            if self.verbose:

                widgets = ['Training : ', Percentage(), ' ', Bar(marker=RotatingMarker()),

                   ' ', ETA(), ' ']

                self.pbar = ProgressBar(widgets=widgets, maxval=(self.max_evals * len(self.param_space)))

                self.pbar.start()

            objective = lambda w: -np.mean([roc_auc_score(y[:, col],

                                            self.calcMean(X[:, col::6], w, training=True))

                                            for col in range(6)])

            self.best_params = fmin(objective, self.param_space, algo=tpe.suggest,

                                    max_evals=self.max_evals)

            if self.verbose:

                print(self.best_params)

        else:

            self.best_params = None

    def predict_proba(self, X):

        """Get predictions."""

        return np.c_[[self.calcMean(X[:, col::6], self.best_params)

                      for col in range(6)]].transpose()

    def calcMean(self, X, w, training = False):

        """Calculate Mean according to weights."""

        self.count += 1

        if self.verbose and self.count <= (self.max_evals * len(self.param_space)) and not self.count%10 and training:

            self.pbar.update(self.count)

        if self.mean == 'simple':

            return np.sum(X, axis=1)/X.shape[1]

        else:

            w = [w[k] for k in sorted(self.sorting, key=self.sorting.get)]

            if self.mean == 'arithmetic':

                return np.sum(X * w, axis=1)/np.sum(w)

            elif self.mean == 'geometric':

                return np.exp(np.sum(np.log(X) * w, axis=1)/np.sum(w))

            elif self.mean == 'power':

                return 1/(1+np.exp(-np.sum(X ** w, axis=1)))

            else:

                print 'Mean should be either "simple", "arithmetic", "geometric" or "power"'