import sklearn
from keras import callbacks

class RocAucCallback(callbacks.Callback):
    def __init__(self,training_data,validation_data):
        self.x = training_data[0]
        self.y = training_data[1]
        self.x_val = validation_data[0]
        self.y_val = validation_data[1]

    def on_train_begin(self, logs={}):
        return self

    def on_train_end(self, logs={}):
        return self

    def on_epoch_begin(self, epoch, logs={}):
        return self

    def on_epoch_end(self, epoch, logs={}):
        #         y_pred = self.model.predict_proba(self.x, verbose=0)
        #         roc = sklearn.metrics.roc_auc_score(self.y, y_pred)
        #         logs['roc_auc'] = sklearn.metrics.roc_auc_score(self.y, y_pred)
        #         logs['norm_gini'] = ( sklearn_metrics.roc_auc_score(self.y, y_pred) * 2 ) - 1

        y_pred_val = self.model.predict_proba(self.x_val, verbose=0)
        roc_val = sklearn.metrics.roc_auc_score(self.y_val, y_pred_val)
        logs['roc_auc_val'] = sklearn.metrics.roc_auc_score(self.y_val, y_pred_val)
        logs['norm_gini_val'] = ( sklearn.metrics.roc_auc_score(self.y_val, y_pred_val) * 2 ) - 1

        #print('\rroc_auc: %s - roc_auc_val: %s - norm_gini: %s - norm_gini_val: %s' % (str(round(roc,5)),str(round(roc_val,5)),str(round((roc*2-1),5)),str(round((roc_val*2-1),5))), end=10*' '+'\n')
        print('roc_auc_val: {}-  norm_gini_val: {}'.format(str(round(roc_val,5)),str(round((roc_val*2-1),5))), end=10*' '+'\n')
        return self

    def on_batch_begin(self, batch, logs={}):
        return self

    def on_batch_end(self, batch, logs={}):
        return self