import torch


def compute_stat_score(pred, label, n, do_argmax: bool = True):
    """Compute metric statistics (true positive, false positive, true negative,
    false negative)
    Parameters
    ----------
    pred : torch.Tensor
        network predictions with probability for each class
    label : torch.Tensor
        ground truth annotation. Classes are given by numeric value
    n : int
        compute scores for class n
    Returns
    -------
    int
        number of true positives
    int
        number of false positives
    int
        number of true negatives
    int
        number of false negatives
    """
    if do_argmax:
        cls_pred = torch.argmax(pred, dim=1)
    else:
        cls_pred = pred[:, 0] > 0.5

    tp = ((cls_pred == n) * (label == n)).to(torch.long).sum()
    fp = ((cls_pred == n) * (label != n)).to(torch.long).sum()
    tn = ((cls_pred != n) * (label != n)).to(torch.long).sum()
    fn = ((cls_pred != n) * (label == n)).to(torch.long).sum()

    return tp, fp, tn, fn