import torch


def calculate_segmentation_statistics(outputs: torch.Tensor, targets: torch.Tensor, class_dim: int = 1, threshold=None):
    """Compute calculate segmentation statistics.

    Args:
        outputs: torch.Tensor.
        targets: torch.Tensor.
        threshold: threshold for binarization of predictions.
        class_dim: indicates class dimension (K).

    Returns:
        True positives , false positives , false negatives for segmentation task.
    """
    num_dims = len(outputs.shape)

    assert num_dims > 2, "Found only two dimensions, shape should be [bs , C , ...]"  # noqa: S101

    assert outputs.shape == targets.shape, "shape mismatch"  # noqa: S101

    if threshold is not None:
        outputs = (outputs > threshold).float()

    dims = [dim for dim in range(num_dims) if dim != class_dim]

    true_positives = torch.sum(outputs * targets, dim=dims)
    false_positives = torch.sum(outputs * (1 - targets), dim=dims)
    false_negatives = torch.sum(targets * (1 - outputs), dim=dims)

    return true_positives, false_positives, false_negatives