# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/metrics/01_detection.ipynb.

# %% auto 0
__all__ = ['map_mar', 'mean_average_precision_recall', 'mean_average_precision_mean_average_recall', 'MeanAveragePrecision',
           'MeanAverageRecall', 'AveragePrecision', 'AverageRecall']

# %% ../../nbs/metrics/01_detection.ipynb 2
from typing import Literal

import torch
from torchmetrics import Metric

from ..utils.bounding_boxes import _IndexedConfidenceScore, get_tps_fps_fns

# %% ../../nbs/metrics/01_detection.ipynb 6
def mean_average_precision_mean_average_recall(
    pred_bboxes: list[torch.Tensor],
    pred_confidence_scores: list[torch.Tensor],
    target_bboxes: list[torch.Tensor],
    target_classes: list[torch.Tensor],
    iou_thresholds: list[float] = [0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95],
    average_precision_num_points: int = 101,
    min_confidence_threshold: float = 0.0,
    max_bboxes_per_image: int | None = 100,
    return_intermediates: bool = False,
) -> tuple[float, float] | tuple[float, float, dict[float, dict[int, float]], dict[float, dict[int, float]]]:
    """Calculate the COCO mean average precision (mAP) for object detection.

    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, 1+num_classes) containing the
            predicted confidence scores for each class. Note that the first column corresponds to the "no-object" class,
            and bounding boxes which fall in this category are ignored.
        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.
        target_classes: A list of length B containing tensors of shape (NT,) containing the target class labels for the
            objects in the image.
        iou_thresholds: A list of IoU thresholds to use for calculating mAP and mAR.
        average_precision_num_points: Number of points over which to calculate average precision.
        min_confidence_threshold: Minimum confidence probability threshold to consider a prediction.
        max_bboxes_per_image: Maximum number of bounding boxes to consider per image. If more are present, only the top
            `max_bboxes_per_image` boxes based on confidence scores are considered. If set to None, all bounding boxes
            are considered.
        return_intermediates: If True, return intermediate values used to calculate mAP and mAR.

    Returns:
        The mean average precision (mAP) and mean average recall (mAR) across all classes and IoU thresholds for the
        entire dataset.
        If `return_intermediates` is True, also returns two dictionaries containing the average precision and average
        recall for each class at each IoU threshold.
    """
    # Some basic tests
    assert len(pred_bboxes) == len(pred_confidence_scores) == len(target_bboxes) == len(target_classes), (
        f"All input lists must have the same length. Got lengths: {len(pred_bboxes)}, {len(pred_confidence_scores)}, "
        f"{len(target_bboxes)}, {len(target_classes)}."
    )
    assert all(
        pred_confidence_score.ndim == 2 for pred_confidence_score in pred_confidence_scores
    ), "Each prediction confidence score input list element must be a 2D tensor."

    # Set some globaly used variables
    B = len(pred_bboxes)
    num_classes = pred_confidence_scores[0].shape[-1] - 1

    # Continue tests
    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(
        pred_confidence_score.shape[-1] == num_classes + 1 for pred_confidence_score in pred_confidence_scores
    ), "Prediction class probabilities must have shape (NP, 1 + num_classes)."
    assert all(
        target_bbox.shape[0] == target_class.shape[0]
        for target_bbox, target_class in zip(target_bboxes, target_classes)
    ), "Each target must have the same number of bounding boxes."
    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)."
    assert all(
        (target_class >= 0).all() and (target_class <= num_classes).all() for target_class in target_classes
    ), f"Target class labels must be between 0 and {num_classes} inclusive."

    # Split everything based on different classes.
    pred_bboxes_by_class = [[] for _ in range(num_classes)]
    pred_confidences_scores_by_class = [[] for _ in range(num_classes)]
    target_bboxes_by_class = [[] for _ in range(num_classes)]
    for b in range(B):
        pred_classes = torch.argmax(pred_confidence_scores[b], dim=-1)
        # (NP,)

        for c in range(num_classes):
            pred_classes_mask = pred_classes == (c + 1)
            # (NP,)
            target_classes_mask = target_classes[b] == (c + 1)
            # (NT,)

            pred_bboxes_by_class[c].append(pred_bboxes[b][pred_classes_mask])
            pred_confidences_scores_by_class[c].append(pred_confidence_scores[b][pred_classes_mask][:, c + 1])
            # (NP,)

            target_bboxes_by_class[c].append(target_bboxes[b][target_classes_mask])
            # (NT,)

    # Limit number of bounding boxes per image if applicable
    if max_bboxes_per_image is not None:
        for b in range(B):
            _confidence_scores = []
            for c in range(num_classes):
                if pred_bboxes_by_class[c][b].numel() > 0:
                    _confidence_scores.extend(
                        _IndexedConfidenceScore.from_batch(
                            pred_confidences_scores_by_class[c][b], batch_index=b, class_index=c
                        )
                    )
            if len(_confidence_scores) == 0:
                continue

            if len(_confidence_scores) > max_bboxes_per_image:
                _confidence_scores = sorted(_confidence_scores, reverse=True)
                topk_confidences = _confidence_scores[:max_bboxes_per_image]
                for c in range(num_classes):
                    topk_confidences_with_class = [x for x in topk_confidences if x.class_index == c]

                    pred_bboxes_by_class[c][b] = torch.stack(
                        [pred_bboxes_by_class[c][b][x.offset_index] for x in topk_confidences_with_class],
                        dim=0,
                        dtype=pred_bboxes_by_class[c][b].dtype,
                        device=pred_bboxes_by_class[c][b].device,
                    )
                    # (NP', 4) or (NP', 6)
                    pred_confidences_scores_by_class[c][b] = torch.stack(
                        [x.score for x in topk_confidences_with_class],
                        dim=0,
                        dtype=pred_confidences_scores_by_class[c][b].dtype,
                        device=pred_confidences_scores_by_class[c][b].device,
                    )
                    # (NP',)

    # For each IOU threshold, calculate average precision and average recall
    average_precisions = {}
    average_recalls = {}
    for iou_threshold in iou_thresholds:
        # For each class calculate average precision and average recall
        class_average_precisions = {}
        class_average_recalls = {}
        for c in range(num_classes):
            # If no target boxes for this class, skip it
            if all(target_bbox.numel() == 0 for target_bbox in target_bboxes_by_class[c]):
                class_average_precisions[c + 1] = float("nan")
                class_average_recalls[c + 1] = float("nan")
                continue

            _, _, _, intermediate_counts = get_tps_fps_fns(
                pred_bboxes=pred_bboxes_by_class[c],
                pred_confidence_scores=pred_confidences_scores_by_class[c],
                target_bboxes=target_bboxes_by_class[c],
                iou_threshold=iou_threshold,
                matching_method="coco",
                min_confidence_threshold=min_confidence_threshold,
                max_bboxes_per_image=None,  # As this has already been done across classes
                return_intermediate_counts=True,
            )
            intermediate_counts = torch.tensor(intermediate_counts, device=pred_bboxes[0].device, dtype=torch.float32)
            # (NC, 3) where the first column is TP, second is FP and third is FN for each prediction considered
            precisions = (intermediate_counts[:, 0] + 1e-5) / (
                intermediate_counts[:, 0] + intermediate_counts[:, 1] + 1e-5
            )
            recalls = (intermediate_counts[:, 0] + 1e-5) / (
                intermediate_counts[:, 0] + intermediate_counts[:, 2] + 1e-5
            )
            # (NC,) each

            # Precision envelope: P_interp(r) = max_{r' >= r} P(r')
            enveloped_precisions = precisions.clone()
            for i in range(len(enveloped_precisions) - 2, -1, -1):
                if enveloped_precisions[i] < enveloped_precisions[i + 1]:
                    enveloped_precisions[i] = enveloped_precisions[i + 1]

            # Calculate average precision using step-wise interpolation
            recall_samples = torch.linspace(0, 1, average_precision_num_points, device=recalls.device)
            idxs = torch.searchsorted(recalls, recall_samples, side="left")
            valid = idxs < enveloped_precisions.numel()
            enveloped_precisions_at_t = torch.zeros_like(recall_samples)
            enveloped_precisions_at_t[valid] = enveloped_precisions[idxs[valid]]
            class_average_precisions[c + 1] = enveloped_precisions_at_t.mean().item()

            # Calculate average recall i.e. maximum recall achieved at this IoU threshold
            class_average_recalls[c + 1] = recalls.max().item() if recalls.numel() > 0 else 0.0

        average_precisions[iou_threshold] = class_average_precisions
        average_recalls[iou_threshold] = class_average_recalls

    map_metric = torch.nanmean(
        torch.stack([torch.tensor(ap) for iou_aps in average_precisions.values() for ap in iou_aps.values()])
    ).item()
    mar_metric = torch.nanmean(
        torch.stack([torch.tensor(ar) for iou_ars in average_recalls.values() for ar in iou_ars.values()])
    ).item()

    if return_intermediates:
        return map_metric, mar_metric, average_precisions, average_recalls
    return map_metric, mar_metric


# Create aliases
map_mar = mean_average_precision_mean_average_recall
mean_average_precision_recall = mean_average_precision_mean_average_recall

# %% ../../nbs/metrics/01_detection.ipynb 12
class _MeanAveragePrecisionMeanAverageRecallBase(Metric):
    """Calculate the COCO mean average precision (mAP) and mean average recall (mAR) for object detection."""

    is_differentiable: bool = False
    higher_is_better: bool = True
    plot_lower_bound: float = 0.0
    plot_upper_bound: float = 1.0

    def __init__(
        self,
        iou_thresholds: list[float] = [0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95],
        average_precision_num_points: int = 101,
        min_confidence_threshold: float = 0.0,
        max_bboxes_per_image: int | None = 100,
        *args,
        **kwargs
    ):
        """Initialize the MeanAveragePrecisionMeanAverageRecall metric.

        Args:
            num_classes: Number of classes in the dataset.
            iou_thresholds: A list of IoU thresholds to use for calculating mAP and mAR.
            average_precision_num_points: Number of points over which to calculate average precision.
            min_confidence_threshold: Minimum confidence score threshold to consider a prediction.
            max_bboxes_per_image: Maximum number of bounding boxes to consider per image. If more are present, only the
                top `max_bboxes_per_image` boxes based on confidence scores are considered.
        """
        super().__init__(*args, **kwargs)

        self.iou_thresholds = iou_thresholds
        self.average_precision_num_points = average_precision_num_points
        self.min_confidence_threshold = min_confidence_threshold
        self.max_bboxes_per_image = max_bboxes_per_image

        self.add_state("pred_bboxes", [], dist_reduce_fx=None, persistent=False)
        self.add_state("pred_confidence_scores", [], dist_reduce_fx=None, persistent=False)
        self.add_state("target_bboxes", [], dist_reduce_fx=None, persistent=False)
        self.add_state("target_classes", [], dist_reduce_fx=None, persistent=False)

    def update(
        self,
        pred_bboxes: list[torch.Tensor],
        pred_confidence_scores: list[torch.Tensor],
        target_bboxes: list[torch.Tensor],
        target_classes: list[torch.Tensor],
    ):
        self.pred_bboxes.extend(pred_bboxes)
        self.pred_confidence_scores.extend(pred_confidence_scores)
        self.target_bboxes.extend(target_bboxes)
        self.target_classes.extend(target_classes)

    def run_functional(self, return_metrics: Literal["map_only", "mar_only"]):
        map_metric, mar_metric = mean_average_precision_mean_average_recall(
            self.pred_bboxes,
            self.pred_confidence_scores,
            self.target_bboxes,
            self.target_classes,
            iou_thresholds=self.iou_thresholds,
            average_precision_num_points=self.average_precision_num_points,
            min_confidence_threshold=self.min_confidence_threshold,
            max_bboxes_per_image=self.max_bboxes_per_image,
        )
        if return_metrics == "map_only":
            return torch.tensor(map_metric, device=self.pred_bboxes[0].device)
        elif return_metrics == "mar_only":
            return torch.tensor(mar_metric, device=self.pred_bboxes[0].device)
        raise NotImplementedError('Only "map_only" and "mar_only" are supported.')

# %% ../../nbs/metrics/01_detection.ipynb 13
class MeanAveragePrecision(_MeanAveragePrecisionMeanAverageRecallBase):
    """Calculate the COCO mean average precision (mAP) for object detection."""

    def compute(self):
        return self.run_functional("map_only")

# %% ../../nbs/metrics/01_detection.ipynb 15
class MeanAverageRecall(_MeanAveragePrecisionMeanAverageRecallBase):
    """Calculate the COCO mean average recall (mAR) for object detection."""

    def compute(self):
        return self.run_functional("mar_only")

# %% ../../nbs/metrics/01_detection.ipynb 17
class AveragePrecision(MeanAveragePrecision):
    """Calculate the COCO average precision (AP) for object detection."""

    def __init__(self, iou_threshold: float, *args, **kwargs):
        iou_thresholds = [iou_threshold]
        super().__init__(iou_thresholds=iou_thresholds, *args, **kwargs)

# %% ../../nbs/metrics/01_detection.ipynb 19
class AverageRecall(MeanAverageRecall):
    """Calculate the COCO average recall (AR) for object detection."""

    def __init__(self, iou_threshold: float, *args, **kwargs):
        iou_thresholds = [iou_threshold]
        super().__init__(iou_thresholds=iou_thresholds, *args, **kwargs)
