import torch


def pre_evaluation_index(y_pred, y_true, sigmoid=False):
    """
    Pre-calculating diffusion matrix for mini-batch evaluation
    Return:
        torch.Tensor((tp, all)) for multi-class classification
        torch.Tensor((tp, fp, fn)) for multi-label classification
    """
    if len(y_true.shape) == 1:
        pred = (y_pred.argmax(1) == y_true).int()
        tp = pred.sum()
        fnp = pred.shape[0] - tp
        return torch.tensor((tp, fnp)).float()
    else:
        if sigmoid:
            border = 0.5
        else:
            border = 0
        y_pred[y_pred >= border] = 1
        y_pred[y_pred < border] = 0
        tp = (y_pred * y_true).sum().to(torch.float32)
        # tn = ((1 - y_true) * (1 - y_pred)).sum().to(torch.float32)
        fp = ((1 - y_true) * y_pred).sum().to(torch.float32)
        fn = (y_true * (1 - y_pred)).sum().to(torch.float32)
        return torch.tensor((tp, fp, fn))