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

"""

Created on Sat Aug 15 21:19:51 2015.

@author: rc, alex

"""

import numpy as np

from sklearn.base import BaseEstimator, ClassifierMixin

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

from preprocessing.aux import delay_preds

import xgboost as xgb

class XGB(BaseEstimator, ClassifierMixin):

    """Ensembling with eXtreme Gradient Boosting."""

    def __init__(self, ensemble, n_estimators=100, max_depth=5, subsample=0.7,

                 nthread=12,delay=None,skip=None,subsample_data=1,partsTest=1, jump=None):

        """Init."""

        self.ensemble = ensemble

        self.n_estimators = n_estimators

        self.max_depth = max_depth

        self.subsample = subsample

        self.nthread = nthread

        ### timecourse history parameters ###

        # how many past time samples to include along with the most recent sample

        self.applyPreds = delay is not None and skip is not None    

        # how many past time samples to include along with the most recent sample  

        self.delay = delay

        # subsample above samples

        self.skip = skip

        # here can be set a custom subsampling scheme, it overrides previous params

        self.jump = jump

        # due to RAM limitations testing data has to be split into 'partsTest' parts

        self.partsTest = partsTest

        # subsampling input data as an efficient form of regularization

        self.subsample_data = subsample_data

        # used in bagging to set different starting points when subsampling the data

        self.mdlNr = 0

        self.clf = []

    def fit(self, X, y):

        """Fit."""

        X = X[(self.mdlNr*5 % self.subsample_data)::self.subsample_data]

        y = y[(self.mdlNr*5 % self.subsample_data)::self.subsample_data]

        if self.applyPreds:

            if self.jump is not None:

                X = delay_preds(X, delay=self.delay/self.subsample_data, skip=self.skip/self.subsample_data, jump=self.jump/self.subsample_data)

            else:

                X = delay_preds(X, delay=self.delay/self.subsample_data, skip=self.skip/self.subsample_data)

        self.clf = []

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

           ' ', ETA(), ' ']

        pbar = ProgressBar(widgets=widgets, maxval=6)

        pbar.start()

        # training separate models for each event

        for col in range(6):

            self.clf.append(xgb.XGBClassifier(n_estimators=self.n_estimators,

                                              max_depth=self.max_depth,

                                              subsample=self.subsample,

                                              nthread=self.nthread))

            self.clf[col].fit(X, y[:, col])

            pbar.update(col)

    def _predict_proba(self,X):

        """Predict probability for each event separately, then concatenate results."""

        pred = []

        for col in range(6):

            pred.append(self.clf[col].predict_proba(X)[:, 1])

        pred = np.vstack(pred).transpose()

        return pred

    def predict_proba(self, X):

        """Predict probability."""

        if self.applyPreds:

            p = np.zeros((X.shape[0],6))

            for part in range(self.partsTest):

                start = part*X.shape[0]//self.partsTest-self.delay*(part>0)

                stop = (part+1)*X.shape[0]//self.partsTest

                X_delayed = delay_preds(X[slice(start,stop)], delay=self.delay, skip=self.skip, jump=self.jump)[self.delay*(part>0):]

                start += self.delay*(part>0)

                p[slice(start,stop)] += self._predict_proba(X_delayed)

                X_delayed = None

            return p

        else:

            return self._predict_proba(X)