# -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function

import keras
import numpy as np
from sklearn.metrics import *


class SKLearnMetrics(keras.callbacks.Callback):
    """ SKLearnMetrics computes various classification metrics at the end of a batch.
     Unforunately, doesn't work when used with generators...."""

    def on_train_begin(self, logs={}):
        self.confusion = []
        self.precision = []
        self.recall = []
        self.f1s = []
        self.kappa = []
        self.auc = []

    def on_epoch_end(self, epoch, logs={}):
        score = np.asarray(self.model.predict(self.validation_data[0]))
        predict = np.round(np.asarray(self.model.predict(self.validation_data[0])))
        target = self.validation_data[1]

        self.auc.append(roc_auc_score(target, score))
        self.confusion.append(confusion_matrix(target, predict))
        self.precision.append(precision_score(target, predict))
        self.recall.append(recall_score(target, predict))
        self.f1s.append(f1_score(target, predict))
        self.kappa.append(cohen_kappa_score(target, predict))
        return