"""Implementation for calculating Mean Average Precision (MAP)."""

from __future__ import annotations

import math
from collections import defaultdict
from collections.abc import Sequence
from uuid import UUID

import torch
from pydantic import BaseModel
from torch import Tensor
from torchmetrics.detection import MeanAveragePrecision

from lightly_purple.server.models.bounding_box_annotation import (
    BoundingBoxAnnotationBase,
)


class DetectionMetricsMAP(BaseModel):
    """Response for computing the MAP detection metric."""

    name: str
    map: float
    map_small: float
    map_medium: float
    map_large: float
    mar_1: float
    mar_10: float
    mar_100: float
    mar_small: float
    mar_medium: float
    mar_large: float
    map_50: float
    map_75: float
    map_per_class: dict[str, float] | None = None
    mar_100_per_class: dict[str, float] | None = None
    classes: list[int]

    @property
    def value(self) -> float:
        """Backwards compatibility alias for map."""
        return self.map

    @property
    def per_class(self) -> dict[str, float] | None:
        """Backwards compatibility alias for map_per_class."""
        return self.map_per_class


def calculate_map_metric(  # noqa: C901
    pred_annotations: Sequence[BoundingBoxAnnotationBase],
    gt_annotations: Sequence[BoundingBoxAnnotationBase],
) -> DetectionMetricsMAP:
    """Calculate the Mean Average Precision (MAP) metric."""
    if not gt_annotations or not pred_annotations:
        return DetectionMetricsMAP(
            name="mAP@.5:.95",
            map=0.0,
            map_small=0.0,
            map_medium=0.0,
            map_large=0.0,
            mar_1=0.0,
            mar_10=0.0,
            mar_100=0.0,
            mar_small=0.0,
            mar_medium=0.0,
            mar_large=0.0,
            map_50=-1.0,
            map_75=-1.0,
            map_per_class=None,
            mar_100_per_class=None,
            classes=[],
        )
    uuid_to_int_label: dict[UUID, int] = {}

    # Map sample_id to annotations.
    sample_id_to_annotations = _group_by_sample_id(
        pred_annotations=pred_annotations,
        gt_annotations=gt_annotations,
    )

    # Create the MeanAveragePrecision object.
    map_metric = MeanAveragePrecision(
        box_format="xywh",
        iou_thresholds=None,
        class_metrics=True,
        average="macro",
        backend="faster_coco_eval",
    )

    # Compute MAP updating the object sample by sample.
    for _sample_id, (
        sample_pred_annotations,
        sample_gt_annotations,
    ) in sample_id_to_annotations.items():
        # Convert to torchmetrics format.
        prediction_tm = _convert_to_torchmetrics(
            annotations=sample_pred_annotations,
            is_prediction=True,
            uuid_to_int_label_map=uuid_to_int_label,
        )
        ground_truth_tm = _convert_to_torchmetrics(
            annotations=sample_gt_annotations,
            is_prediction=False,
            uuid_to_int_label_map=uuid_to_int_label,
        )

        # Update the metric.
        map_metric.update(preds=[prediction_tm], target=[ground_truth_tm])

    # Compute the final results.
    results = map_metric.compute()
    # Invert label map to recover original UUIDs
    int_to_uuid = {v: k for k, v in uuid_to_int_label.items()}

    # Helper to convert Tensor to Python list
    def to_list(t: Tensor) -> list[float]:
        if t.dim() == 0:
            return [t.item()]
        return t.tolist()

    # Scalar metrics
    map_val = results["map"].item()
    map_small = results["map_small"].item()
    map_medium = results["map_medium"].item()
    map_large = results["map_large"].item()
    mar_1 = results["mar_1"].item()
    mar_10 = results["mar_10"].item()
    mar_100 = results["mar_100"].item()
    mar_small = results["mar_small"].item()
    mar_medium = results["mar_medium"].item()
    mar_large = results["mar_large"].item()

    # IoU-specific metrics
    map_50 = results.get("map_50")
    map_50 = (
        map_50.item()
        if isinstance(map_50, Tensor) and map_50.dim() == 0
        else -1.0
    )
    map_75 = results.get("map_75")
    map_75 = (
        map_75.item()
        if isinstance(map_75, Tensor) and map_75.dim() == 0
        else -1.0
    )

    # Per-class metrics
    raw_map_pc = results.get("map_per_class")
    per_class_dict: dict[str, float] | None = None
    if raw_map_pc is not None:
        pc_list = to_list(raw_map_pc)
        per_class_dict = {}
        for idx, val in enumerate(pc_list):
            if not math.isnan(val):
                per_class_dict[str(int_to_uuid[idx])] = val

    # Per-class recall at 100 detections
    raw_mar100_pc = results.get("mar_100_per_class")
    mar_100_pc_dict: dict[str, float] | None = None
    if raw_mar100_pc is not None:
        m100_list = to_list(raw_mar100_pc)
        mar_100_pc_dict = {}
        for idx, val in enumerate(m100_list):
            if not math.isnan(val):
                mar_100_pc_dict[str(int_to_uuid[idx])] = val

    # Observed classes
    raw_classes = results.get("classes")
    classes_list_output: list[int] = []
    if raw_classes is not None:
        classes_list_output = [int(x) for x in to_list(raw_classes)]

    return DetectionMetricsMAP(
        name="mAP@.5:.95",
        map=map_val,
        map_small=map_small,
        map_medium=map_medium,
        map_large=map_large,
        mar_1=mar_1,
        mar_10=mar_10,
        mar_100=mar_100,
        mar_small=mar_small,
        mar_medium=mar_medium,
        mar_large=mar_large,
        map_50=map_50,
        map_75=map_75,
        map_per_class=per_class_dict,
        mar_100_per_class=mar_100_pc_dict,
        classes=classes_list_output,
    )


def _group_by_sample_id(
    pred_annotations: Sequence[BoundingBoxAnnotationBase],
    gt_annotations: Sequence[BoundingBoxAnnotationBase],
) -> dict[
    UUID,
    tuple[list[BoundingBoxAnnotationBase], list[BoundingBoxAnnotationBase]],
]:
    """Group prediction and ground truth annotations by sample_id.

    Returns a dictionary with sample_id as key and a tuple of
    (list of prediction annotations, list of ground truth annotations) as value.
    """
    sample_id_to_annotations: defaultdict[
        UUID,
        tuple[list[BoundingBoxAnnotationBase], list[BoundingBoxAnnotationBase]],
    ] = defaultdict(lambda: ([], []))
    for ann in pred_annotations:
        sample_id_to_annotations[ann.sample_id][0].append(ann)
    for ann in gt_annotations:
        sample_id_to_annotations[ann.sample_id][1].append(ann)
    return sample_id_to_annotations


def _convert_to_torchmetrics(
    annotations: Sequence[BoundingBoxAnnotationBase],
    is_prediction: bool,
    uuid_to_int_label_map: dict[UUID, int],
) -> dict[str, Tensor]:
    """Convert annotations to torchmetrics format.

    Args:
        annotations: List of bounding box annotations.
        is_prediction: Whether the annotations are predictions.
        uuid_to_int_label_map: Map from UUID to integer label.

    Returns:
        Dictionary in torchmetrics format. For predictions the keys are
        `boxes`, `scores`, `labels`. For ground truth they are `boxes`,
        `labels`.

    """
    if not annotations:
        empty_boxes = torch.empty((0, 4), dtype=torch.float32)
        empty_labels = torch.empty((0,), dtype=torch.int64)
        if is_prediction:
            return {
                "boxes": empty_boxes,
                "scores": torch.empty((0,), dtype=torch.float32),
                "labels": empty_labels,
            }
        return {"boxes": empty_boxes, "labels": empty_labels}

    boxes = torch.tensor(
        [[a.x, a.y, a.width, a.height] for a in annotations],
        dtype=torch.float32,
    )

    mapped_labels = []
    # Use the passed-in map instead of a global one
    for uuid_label in (a.annotation_label_id for a in annotations):
        if uuid_label not in uuid_to_int_label_map:
            # Assign the next available integer ID based on the current map size
            uuid_to_int_label_map[uuid_label] = len(uuid_to_int_label_map)
        mapped_labels.append(uuid_to_int_label_map[uuid_label])

    labels = torch.tensor(mapped_labels, dtype=torch.int64)

    if is_prediction:
        scores = torch.tensor(
            [a.confidence for a in annotations], dtype=torch.float32
        )
        return {
            "boxes": boxes,
            "scores": scores,
            "labels": labels,
        }
    return {"boxes": boxes, "labels": labels}
