import torch
import numpy as np
from typing import Optional


@torch.no_grad()
def compute_dice(pred: torch.Tensor, target: torch.Tensor, num_classes: int) -> float:
    """
    Compute mean Dice coefficient (F1 score) for a batch of segmentation predictions.

    The Dice coefficient for a class is defined as:
        Dice = 2 * |X ∩ Y| / (|X| + |Y|)
    where X is the predicted binary mask for that class, and Y is the ground truth mask.
    This function averages the per-class Dice scores over all classes that appear in the batch.

    Args:
        pred (torch.Tensor): Raw logits or class scores of shape (N, C, H, W),
            where C is the number of classes. Argmax is applied along the class dimension.
        target (torch.Tensor): Ground truth class indices of shape (N, H, W).
            Values 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 Dice coefficient over all classes that have at least one pixel
            in either the prediction or the ground truth for the batch.
            Returns 0.0 if no class appears (i.e., all masks are empty).

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

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

        intersection = (pred_mask & target_mask).sum().float()
        cardinality = pred_mask.sum().float() + target_mask.sum().float()

        if cardinality > 0:
            dice_list.append((2 * intersection / cardinality).item())
        # If cardinality == 0, the class is not present in either mask; skip it

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