# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/utils/11_bounding_boxes.ipynb.

# %% auto 0
__all__ = ['get_tps_fps_fns']

# %% ../../nbs/utils/11_bounding_boxes.ipynb 2
from typing import Literal

import torch
from monai.data.box_utils import box_iou
from scipy.optimize import linear_sum_assignment

# %% ../../nbs/utils/11_bounding_boxes.ipynb 6
class _IndexedConfidenceScore:
    def __init__(self, confidence: float, batch_index: int, offset_index: int, class_index: int | None = None):
        self.confidence = confidence
        self.batch_index = batch_index
        self.offset_index = offset_index
        self.class_index = class_index

    @classmethod
    def from_batch(cls, confidence_scores: torch.Tensor, batch_index: int, class_index: int | None = None):
        return [
            cls(
                confidence=confidence_scores[i].detach().cpu().item(),
                batch_index=batch_index,
                offset_index=i,
                class_index=class_index,
            )
            for i in range(len(confidence_scores))
        ]

    def __gt__(self, other: "_IndexedConfidenceScore"):
        return self.confidence > other.confidence

    def __lt__(self, other: "_IndexedConfidenceScore"):
        return self.confidence < other.confidence

    def __ge__(self, other: "_IndexedConfidenceScore"):
        return self.confidence >= other.confidence

    def __le__(self, other: "_IndexedConfidenceScore"):
        return self.confidence <= other.confidence

    def __repr__(self):
        _repr = f"_IndexedConfidenceScore(confidence={self.confidence}, batch_index={self.batch_index}, offset_index={self.offset_index})"
        if self.class_index is not None:
            _repr += f", class_index={self.class_index}"
        return _repr

# %% ../../nbs/utils/11_bounding_boxes.ipynb 9
def get_tps_fps_fns(
    pred_bboxes: list[torch.Tensor],
    pred_confidence_scores: list[torch.Tensor],
    target_bboxes: list[torch.Tensor],
    iou_threshold: float,
    matching_method: Literal["coco", "hungarian"] = "coco",
    min_confidence_threshold: float = 0.0,
    max_bboxes_per_image: int | None = None,
    return_intermediate_counts: bool = False,
) -> (
    tuple[set[tuple[int, int, int]], set[tuple[int, int]], set[tuple[int, int]]]
    | tuple[set, set, set, list[tuple[int, int, int]]]
):
    """Given predicted and target bounding boxes, their confidence scores, and an IOU threshold, get a matching of
    true positives, and a set of false positives and false negatives.

    Args:
        pred_bboxes: A list of length B containing tensors of shape (NP, 4) or (NP, 6) containing the predicted bounding
            box parameters in xyxy or xyzxyz format.
        pred_confidence_scores: A list of length B containing tensors of shape (NP,) containing the predicted confidence
            scores for the corresponding bounding boxes.
        target_bboxes: A list of length B containing tensors of shape (NT, 4) or (NT, 6) containing the target bounding
            box parameters in xyxy or xyzxyz format.
        iou_threshold: The IOU threshold above which a predicted box is considered a match for a target box.
        matching_method: The method to use for matching predicted boxes to target boxes. 'coco' implements the greedy
            matching algorithm used in the COCO dataset. 'hungarian' implements the Hungarian algorithm for optimal
            matching. Note that 'hungarian' is more computationally expensive and may not scale well to large numbers of
            boxes.
        min_confidence_threshold: Minimum confidence score for a predicted box to be considered for matching.
        max_bboxes_per_image: If not None, consider only the top K predicted boxes per image based on confidence scores.
        return_intermediate_counts: Whether to return intermediate counts of true positives, false positives and false
            negatives after each prediction is considered. Useful for plotting precision-recall curves.

    Returns:
        The first set contains tuples of (b, p, t) where b is the batch index, p is the index of the predicted box
        and t is the index of the matched target box. The second set contains tuples of (b, p) where b is the batch
        index and p is the index of the false positive predicted box. The third set contains tuples of (b, t) where b
        is the batch index and t is the index of the false negative target box.
        If `return_intermediate_counts` is True, also returns a list of tuples of (TP, FP, FN) counts after each
        prediction.
    """
    assert (
        len(pred_bboxes) == len(pred_confidence_scores) == len(target_bboxes)
    ), "Batch size must be the same for all inputs"
    assert matching_method in ["coco", "hungarian"], "matching_method must be either 'coco' or 'hungarian'"
    assert all(
        pred_bbox.shape[0] == pred_confidence_score.shape[0]
        for pred_bbox, pred_confidence_score in zip(pred_bboxes, pred_confidence_scores)
    ), "Each prediction input list element must have the same number of bounding boxes."
    assert all(
        pred_bbox.shape[1] == 4 or pred_bbox.shape[1] == 6 for pred_bbox in pred_bboxes
    ), "Prediction bounding boxes must have shape (NP, 4) or (NP, 6)."
    assert all(
        target_bbox.shape[1] == 4 or target_bbox.shape[1] == 6 for target_bbox in target_bboxes
    ), "Target bounding boxes must have shape (NT, 4) or (NT, 6)."

    B = len(pred_bboxes)

    # Join all confidence scores and keep track of which batch and which box they correspond to
    all_pred_confidence_scores = []
    for b in range(B):
        _confidence_scores = _IndexedConfidenceScore.from_batch(pred_confidence_scores[b], batch_index=b)
        if max_bboxes_per_image is not None and len(_confidence_scores) > max_bboxes_per_image:
            _confidence_scores = sorted(_confidence_scores, reverse=True)[:max_bboxes_per_image]
        all_pred_confidence_scores.extend(_confidence_scores)
    all_pred_confidence_scores = sorted(all_pred_confidence_scores, reverse=True)

    # Calculate IOUs between all predicted and target boxes
    ious = []
    hungarian_matchings = []
    for b in range(B):
        _ious = box_iou(pred_bboxes[b], target_bboxes[b])
        ious.append(_ious)
        # (NP, NT), where NP is number of predicted boxes and NT is number of target boxes

        if matching_method == "hungarian":
            if _ious.numel() > 0:
                # Calculate optimal matching using Hungarian algorithm
                _matching = linear_sum_assignment(-_ious.cpu().numpy())
                hungarian_matchings.append(_matching)
            else:
                hungarian_matchings.append((torch.tensor([], dtype=torch.int64), torch.tensor([], dtype=torch.int64)))

    matched_target_indices = [set() for _ in range(B)]
    tps, fps, fns = set(), set(), {(b, i) for b in range(B) for i in range(len(target_bboxes[b]))}

    # In descending order, update tp, fp, fn and calculate precision and recall at each step
    intermediate_counts = []
    for indexed_confidence_score in all_pred_confidence_scores:
        confidence_score, b, pred_offset = (
            indexed_confidence_score.confidence,
            indexed_confidence_score.batch_index,
            indexed_confidence_score.offset_index,
        )

        if confidence_score < min_confidence_threshold:
            # Do not consider this and following predictions if confidence score is below threshold
            break

        if matching_method == "coco":
            # COCO-style greedy matching
            pred_ious: torch.Tensor = ious[b][pred_offset].clone()
            # (NT,)

            if pred_ious.numel() > 0:
                if matched_target_indices[b]:
                    # Exclude already matched target boxes
                    pred_ious[list(matched_target_indices[b])] = -1.0
                # Exclude target boxes below IOU threshold
                pred_ious[pred_ious < iou_threshold] = -1.0

            if pred_ious.numel() == 0 or pred_ious.amax() < 0.0:
                # No valid target box to match with
                fps.add((b, pred_offset))
            else:
                target_offset = pred_ious.argmax().item()
                matched_target_indices[b].add(target_offset)
                tps.add((b, pred_offset, target_offset))
                fns.discard((b, target_offset))
        else:
            matched_pred_offsets, matched_target_offsets = hungarian_matchings[b]
            if pred_offset in matched_pred_offsets:
                target_offset_index = (matched_pred_offsets == pred_offset).nonzero()[0].item()
                target_offset = matched_target_offsets[target_offset_index]
                if ious[b][pred_offset, target_offset] >= iou_threshold:
                    tps.add((b, pred_offset, target_offset))
                    fns.discard((b, target_offset))
                else:
                    fps.add((b, pred_offset))
            else:
                fps.add((b, pred_offset))

        if return_intermediate_counts:
            intermediate_counts.append((len(tps), len(fps), len(fns)))

    if return_intermediate_counts:
        return tps, fps, fns, intermediate_counts
    return tps, fps, fns
