import torch
import numpy as np
from typing import Optional


@torch.no_grad()
def compute_iou(pred: torch.Tensor, target: torch.Tensor, num_classes: int) -> float:
    """
    Compute mean Intersection-over-Union (mIoU) for a batch of segmentation predictions.

    The function takes raw logits or class probabilities, converts them to class predictions
    via argmax, and calculates the IoU for each class present in the batch. The final result
    is the average of these per-class IoUs.

    Args:
        pred (torch.Tensor): Raw logits or class scores of shape (N, C, H, W),
            where C is the number of classes. The values are typically not normalized,
            as argmax is applied along the class dimension.
        target (torch.Tensor): Ground truth class indices of shape (N, H, W).
            Each value should be in the range [0, C-1].
        num_classes (int): Total number of classes. Used to iterate over all possible classes.

    Returns:
        float: Mean IoU over all classes that have at least one positive pixel in the batch.
            Returns 0.0 if no class has any positive pixel (i.e., union == 0 for all classes).

    Examples:
        >>> pred = torch.randn(2, 10, 128, 128)   # batch of 2, 10 classes
        >>> target = torch.randint(0, 10, (2, 128, 128))
        >>> miou = compute_iou(pred, target, num_classes=10)
        >>> print(f"mIoU: {miou:.4f}")
    """
    # Convert logits to predicted class indices
    pred = torch.argmax(pred, dim=1)  # (N, H, W)

    iou_list = []
    for cls in range(num_classes):
        pred_mask = (pred == cls)
        target_mask = (target == cls)

        intersection = (pred_mask & target_mask).sum().float()
        union = (pred_mask | target_mask).sum().float()

        if union > 0:
            iou_list.append((intersection / union).item())
        # If union == 0, the class is not present in the batch; skip it

    # Return the mean of all present classes, or 0.0 if none were present
    return float(np.mean(iou_list)) if iou_list else 0.0