import torch
from typing import Tuple


def accuracy(preds: torch.Tensor, targets: torch.Tensor, cls_dim: int = 1,
             ) -> Tuple[float, float, float, float]:
    """Calculate accuracy for given predictions and one-hot-encoded binary targets.
    The class dimension is assumed to be in dimension cls_dim.
    If the size of cls_dim is 1 in the predictions, it is assumed to be a binary classification
    output in one node (non-one-hot-encoded) and the values are sigmoided and thresholded by 0.5 to
    obtain the class prediction.

    :param preds: predictions tensor
    :param targets: targets tensor
    :param cls_dim: the dimension in which to find the class information in case of multi-class
        classification; e.g. if cls_dim == 1 the prediction tensors are assumed to be of shape
        (batch, num_classes, *)
    :return: tuple of (accuracy, #true pos, #true neg, total number of samples)
    """
    # Calc true positives, true negatives, and tn+tp+fn+fp for batch
    # from one-hot or binary encoding to class indices:
    class_pred = torch.argmax(preds, dim=cls_dim) if preds.size()[cls_dim] > 1 else \
        (torch.sigmoid(preds) > 0.5).squeeze(cls_dim).float()

    class_gt = torch.argmax(targets, dim=cls_dim).float()
    # accuracy calculation
    batch_tp = float(torch.sum(class_pred * class_gt))
    batch_tn = float(torch.sum((1 - class_pred) * (1 - class_gt)))
    batch_all = float(class_gt.size()[0])
    batch_acc = (batch_tp + batch_tn) / batch_all
    return batch_acc, batch_tp, batch_tn, batch_all