import numpy as np


class Evaluator(object):
    def __init__(self, num_class):
        self.num_class = num_class
        self.confusion_matrix = np.zeros((self.num_class,) * 2)

    def Acc(self, session_class):
        confusion_matrix = self.confusion_matrix[:session_class, :session_class]
        return confusion_matrix, np.diag(confusion_matrix) / np.sum(confusion_matrix, axis=1), np.sum(np.diag(confusion_matrix)) / np.sum(confusion_matrix)
        # return np.diag(confusion_matrix) / np.sum(confusion_matrix, axis=1)

    def _generate_matrix(self, gt_image, pre_image):
        mask = (gt_image >= 0) & (gt_image < self.num_class)
        label = self.num_class * gt_image[mask].astype('int') + pre_image[mask]
        count = np.bincount(label, minlength=self.num_class ** 2)
        confusion_matrix = count.reshape(self.num_class, self.num_class)
        return confusion_matrix

    def add_batch(self, gt_image, pre_image):
        self.confusion_matrix += self._generate_matrix(gt_image, pre_image)

    def reset(self):
        self.confusion_matrix = np.zeros((self.num_class,) * 2)
