from typing import List
from typing import Dict
import torch


def _get_top_segmentation_masks(
    batch_segmentation: List[Dict[str, torch.Tensor]], class_index: int
) -> List[torch.Tensor]:
    """
    For a segmentation output by a segmentation model, find the top-scoring mask for
    each image associated with the class_index class.

    Parameters
    ----------
    batch_segmentation: The output from the segmentation model.
    class_index: The class to pick out masks from.

    Returns
    -------
    top_masks: Tensor of segmentation masks of shape (B, C, H, W) where C is a single
        dimension (I think!).
    """
    top_masks = []
    for segmentation in batch_segmentation:
        labels = segmentation["labels"].to("cpu")
        masks = segmentation["masks"].to("cpu")
        scores = segmentation["scores"].to("cpu")
        best_score = -1
        best_mask = None
        for label, mask, score in zip(labels, masks, scores):
            if label == class_index:
                if score > best_score:
                    best_mask = mask
                    best_score = score
        if best_mask is None:
            raise ValueError(f"No mask found for class index {class_index}")
        # Unsqueeze 0 dim because we lost that dimension while iterating through the for
        # loop.
        top_masks.append(best_mask.unsqueeze(0))

    return top_masks