"""Reusable object-detection metrics based on IoU matching."""

from __future__ import annotations

from typing import TYPE_CHECKING

import torch


if TYPE_CHECKING:
    from collections.abc import Iterable, Mapping, MutableMapping, Sequence


_EPSILON = 1e-6


def _to_tensor(values: torch.Tensor | Sequence[float] | Sequence[int], *, dtype: torch.dtype) -> torch.Tensor:
    """Convert arbitrary tensor-like values into a detached tensor.

    Args:
        values: Tensor-like values to convert.
        dtype: Target dtype for resulting tensor.

    Returns:
        Detached tensor with requested dtype.
    """
    if isinstance(values, torch.Tensor):
        return values.detach().to(dtype=dtype)
    return torch.as_tensor(values, dtype=dtype)


def _extract_detection_item(record: Mapping[str, object], key: str, default: torch.Tensor) -> torch.Tensor:
    """Extract tensor field from detection dictionary with fallback.

    Args:
        record: Detection record mapping containing tensor-like entries.
        key: Key to read from record.
        default: Default tensor returned when key is missing.

    Returns:
        Detached tensor for key value, or provided default tensor.
    """
    value = record.get(key)
    if value is None:
        return default
    if isinstance(value, torch.Tensor):
        return value.detach()
    return torch.as_tensor(value)


def _class_ids(num_classes: int, background_class: int | None) -> list[int]:
    """Return class ids to evaluate while optionally skipping background.

    Args:
        num_classes: Total number of classes in evaluation.
        background_class: Optional class id to exclude from evaluation.

    Returns:
        List of class ids included in metric computation.
    """
    ids = list(range(num_classes))
    if background_class is not None and 0 <= background_class < num_classes:
        ids.remove(background_class)
    return ids


def box_iou(boxes1: torch.Tensor, boxes2: torch.Tensor) -> torch.Tensor:
    """Compute IoU matrix for ``xyxy`` boxes.

    Args:
        boxes1: First box tensor with shape ``[N, 4]``.
        boxes2: Second box tensor with shape ``[M, 4]``.

    Returns:
        IoU matrix tensor with shape ``[N, M]``.
    """
    boxes1 = _to_tensor(boxes1, dtype=torch.float32).reshape(-1, 4)
    boxes2 = _to_tensor(boxes2, dtype=torch.float32).reshape(-1, 4)

    if boxes1.numel() == 0 or boxes2.numel() == 0:
        return torch.zeros((boxes1.size(0), boxes2.size(0)), dtype=torch.float32)

    top_left = torch.maximum(boxes1[:, None, :2], boxes2[None, :, :2])
    bottom_right = torch.minimum(boxes1[:, None, 2:], boxes2[None, :, 2:])
    wh = (bottom_right - top_left).clamp(min=0)
    intersection = wh[..., 0] * wh[..., 1]

    area1 = ((boxes1[:, 2] - boxes1[:, 0]).clamp(min=0) * (boxes1[:, 3] - boxes1[:, 1]).clamp(min=0))
    area2 = ((boxes2[:, 2] - boxes2[:, 0]).clamp(min=0) * (boxes2[:, 3] - boxes2[:, 1]).clamp(min=0))
    union = area1[:, None] + area2[None, :] - intersection

    return intersection / union.clamp(min=_EPSILON)


def _match_single_image(
    pred_boxes: torch.Tensor,
    pred_scores: torch.Tensor,
    target_boxes: torch.Tensor,
    iou_threshold: float,
) -> tuple[int, int, int, list[tuple[float, int]]]:
    """Perform greedy IoU matching for one class in one image.

    Args:
        pred_boxes: Predicted boxes tensor for a single class.
        pred_scores: Confidence scores for predicted boxes.
        target_boxes: Ground-truth boxes tensor for the same class.
        iou_threshold: IoU threshold for true-positive matching.

    Returns:
        Tuple ``(tp, fp, fn, ranked_predictions)`` for the image/class pair.
    """
    pred_boxes = pred_boxes.reshape(-1, 4)
    pred_scores = pred_scores.reshape(-1)
    target_boxes = target_boxes.reshape(-1, 4)

    if pred_boxes.numel() == 0 and target_boxes.numel() == 0:
        return 0, 0, 0, []
    if pred_boxes.numel() == 0:
        fn = int(target_boxes.size(0))
        return 0, 0, fn, []
    if target_boxes.numel() == 0:
        fp = int(pred_boxes.size(0))
        ranked = [(float(score.item()), 0) for score in pred_scores]
        return 0, fp, 0, ranked

    order = torch.argsort(pred_scores, descending=True)
    pred_boxes = pred_boxes[order]
    pred_scores = pred_scores[order]
    ious = box_iou(pred_boxes, target_boxes)

    matched_gt = torch.zeros(target_boxes.size(0), dtype=torch.bool)
    tp = 0
    fp = 0
    ranked_predictions: list[tuple[float, int]] = []

    for pred_idx in range(pred_boxes.size(0)):
        iou_row = ious[pred_idx]
        best_iou, best_gt = torch.max(iou_row, dim=0)
        is_true_positive = bool(best_iou.item() >= iou_threshold and not matched_gt[best_gt].item())
        if is_true_positive:
            matched_gt[best_gt] = True
            tp += 1
            ranked_predictions.append((float(pred_scores[pred_idx].item()), 1))
        else:
            fp += 1
            ranked_predictions.append((float(pred_scores[pred_idx].item()), 0))

    fn = int((~matched_gt).sum().item())
    return tp, fp, fn, ranked_predictions


def compute_metrics(
    outputs: Iterable[Mapping[str, object]],
    targets: Iterable[Mapping[str, object]],
    num_classes: int,
    iou_threshold: float = 0.5,
    background_class: int | None = 0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Compute per-class precision, recall, and F1-score for detections.

    Args:
        outputs: Iterable of model output dictionaries.
        targets: Iterable of target annotation dictionaries.
        num_classes: Total number of classes in evaluation.
        iou_threshold: IoU threshold used for matching.
        background_class: Optional class id to exclude from evaluation.

    Returns:
        Tuple of tensors ``(precision, recall, f1_score)`` per class.
    """
    all_true_positives = torch.zeros(num_classes, dtype=torch.float32)
    all_false_positives = torch.zeros(num_classes, dtype=torch.float32)
    all_false_negatives = torch.zeros(num_classes, dtype=torch.float32)

    classes = _class_ids(num_classes=num_classes, background_class=background_class)
    for output, target in zip(outputs, targets, strict=False):
        output_boxes = _extract_detection_item(output, "boxes", torch.zeros((0, 4), dtype=torch.float32)).to(
            dtype=torch.float32, device="cpu"
        )
        output_labels = _extract_detection_item(output, "labels", torch.zeros((0,), dtype=torch.long)).to(
            dtype=torch.long, device="cpu"
        )
        output_scores = _extract_detection_item(
            output,
            "scores",
            torch.ones((output_boxes.size(0),), dtype=torch.float32),
        ).to(dtype=torch.float32, device="cpu")
        target_boxes = _extract_detection_item(target, "boxes", torch.zeros((0, 4), dtype=torch.float32)).to(
            dtype=torch.float32, device="cpu"
        )
        target_labels = _extract_detection_item(target, "labels", torch.zeros((0,), dtype=torch.long)).to(
            dtype=torch.long, device="cpu"
        )

        for class_id in classes:
            pred_mask = output_labels == class_id
            gt_mask = target_labels == class_id

            class_pred_boxes = output_boxes[pred_mask]
            class_pred_scores = output_scores[pred_mask]
            class_target_boxes = target_boxes[gt_mask]

            tp, fp, fn, _ = _match_single_image(
                pred_boxes=class_pred_boxes,
                pred_scores=class_pred_scores,
                target_boxes=class_target_boxes,
                iou_threshold=iou_threshold,
            )

            all_true_positives[class_id] += float(tp)
            all_false_positives[class_id] += float(fp)
            all_false_negatives[class_id] += float(fn)

    precision = all_true_positives / (all_true_positives + all_false_positives + _EPSILON)
    recall = all_true_positives / (all_true_positives + all_false_negatives + _EPSILON)
    f1_score = 2.0 * precision * recall / (precision + recall + _EPSILON)

    return precision, recall, f1_score


def _compute_ap(recall: torch.Tensor, precision: torch.Tensor) -> float:
    """Compute average precision from precision-recall curve.

    Args:
        recall: Monotonic recall tensor values.
        precision: Precision tensor values paired with recall.

    Returns:
        Average precision scalar.
    """
    if recall.numel() == 0 or precision.numel() == 0:
        return 0.0

    mrec = torch.cat(
        [torch.tensor([0.0], dtype=torch.float32), recall.to(torch.float32), torch.tensor([1.0], dtype=torch.float32)]
    )
    mpre = torch.cat(
        [torch.tensor([0.0], dtype=torch.float32), precision.to(torch.float32), torch.tensor([0.0], dtype=torch.float32)]
    )

    for i in range(mpre.numel() - 1, 0, -1):
        mpre[i - 1] = torch.maximum(mpre[i - 1], mpre[i])

    changing_points = torch.nonzero(mrec[1:] != mrec[:-1], as_tuple=False).reshape(-1)
    if changing_points.numel() == 0:
        return 0.0

    ap = torch.sum((mrec[changing_points + 1] - mrec[changing_points]) * mpre[changing_points + 1])
    return float(ap.item())


def _average_precision_per_class(
    outputs: Iterable[Mapping[str, object]],
    targets: Iterable[Mapping[str, object]],
    num_classes: int,
    iou_threshold: float,
    background_class: int | None,
) -> tuple[dict[int, float], list[int]]:
    """Compute AP per class at a single IoU threshold.

    Args:
        outputs: Iterable of model output dictionaries.
        targets: Iterable of target annotation dictionaries.
        num_classes: Total number of classes in evaluation.
        iou_threshold: IoU threshold used for matching.
        background_class: Optional class id to exclude from evaluation.

    Returns:
        Tuple ``(ap_by_class, valid_classes)`` for the threshold.
    """
    class_ids = _class_ids(num_classes=num_classes, background_class=background_class)
    score_matches: dict[int, list[tuple[float, int]]] = {class_id: [] for class_id in class_ids}
    total_gt: dict[int, int] = dict.fromkeys(class_ids, 0)

    for output, target in zip(outputs, targets, strict=False):
        output_boxes = _extract_detection_item(output, "boxes", torch.zeros((0, 4), dtype=torch.float32)).to(
            dtype=torch.float32, device="cpu"
        )
        output_labels = _extract_detection_item(output, "labels", torch.zeros((0,), dtype=torch.long)).to(
            dtype=torch.long, device="cpu"
        )
        output_scores = _extract_detection_item(
            output,
            "scores",
            torch.ones((output_boxes.size(0),), dtype=torch.float32),
        ).to(dtype=torch.float32, device="cpu")
        target_boxes = _extract_detection_item(target, "boxes", torch.zeros((0, 4), dtype=torch.float32)).to(
            dtype=torch.float32, device="cpu"
        )
        target_labels = _extract_detection_item(target, "labels", torch.zeros((0,), dtype=torch.long)).to(
            dtype=torch.long, device="cpu"
        )

        for class_id in class_ids:
            pred_mask = output_labels == class_id
            gt_mask = target_labels == class_id

            class_pred_boxes = output_boxes[pred_mask]
            class_pred_scores = output_scores[pred_mask]
            class_target_boxes = target_boxes[gt_mask]
            total_gt[class_id] += int(class_target_boxes.size(0))

            _, _, _, ranked = _match_single_image(
                pred_boxes=class_pred_boxes,
                pred_scores=class_pred_scores,
                target_boxes=class_target_boxes,
                iou_threshold=iou_threshold,
            )
            score_matches[class_id].extend(ranked)

    ap_by_class: dict[int, float] = {}
    valid_classes: list[int] = []
    for class_id in class_ids:
        if total_gt[class_id] <= 0:
            ap_by_class[class_id] = 0.0
            continue

        valid_classes.append(class_id)
        ranked = sorted(score_matches[class_id], key=lambda item: item[0], reverse=True)
        if not ranked:
            ap_by_class[class_id] = 0.0
            continue

        tp = torch.tensor([is_tp for _, is_tp in ranked], dtype=torch.float32)
        fp = 1.0 - tp
        tp_cum = torch.cumsum(tp, dim=0)
        fp_cum = torch.cumsum(fp, dim=0)
        precision = tp_cum / (tp_cum + fp_cum + _EPSILON)
        recall = tp_cum / (float(total_gt[class_id]) + _EPSILON)
        ap_by_class[class_id] = _compute_ap(recall=recall, precision=precision)

    return ap_by_class, valid_classes


def compute_map_metrics(
    outputs: Iterable[Mapping[str, object]],
    targets: Iterable[Mapping[str, object]],
    num_classes: int,
    iou_thresholds: Sequence[float] = (0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95),
    background_class: int | None = 0,
) -> tuple[float, float, dict[int, float]]:
    """Compute ``mAP@50``, ``mAP@50-95``, and per-class ``AP@50``.

    Args:
        outputs: Iterable of model output dictionaries.
        targets: Iterable of target annotation dictionaries.
        num_classes: Total number of classes in evaluation.
        iou_thresholds: Sequence of IoU thresholds for mAP computation.
        background_class: Optional class id to exclude from evaluation.

    Returns:
        Tuple ``(map50, map50_95, per_class_ap_50)``.
    """
    thresholds = list(dict.fromkeys(float(threshold) for threshold in iou_thresholds))
    if 0.5 not in thresholds:
        thresholds.insert(0, 0.5)

    map_by_threshold: dict[float, float] = {}
    per_class_ap_50: dict[int, float] = {}

    cached_outputs = list(outputs)
    cached_targets = list(targets)

    for threshold in thresholds:
        ap_by_class, valid_classes = _average_precision_per_class(
            outputs=cached_outputs,
            targets=cached_targets,
            num_classes=num_classes,
            iou_threshold=threshold,
            background_class=background_class,
        )
        if threshold == 0.5:
            per_class_ap_50 = ap_by_class

        if not valid_classes:
            map_by_threshold[threshold] = 0.0
        else:
            map_by_threshold[threshold] = float(sum(ap_by_class[class_id] for class_id in valid_classes) / len(valid_classes))

    map50 = map_by_threshold.get(0.5, 0.0)
    map50_95 = float(sum(map_by_threshold.values()) / len(map_by_threshold)) if map_by_threshold else 0.0
    return map50, map50_95, per_class_ap_50


def epoch_metrics(
    data_loader: Iterable[tuple[Sequence[torch.Tensor], Sequence[Mapping[str, object]]]],
    model: torch.nn.Module,
    device: torch.device,
    num_classes: int,
    iou_threshold: float = 0.5,
    background_class: int | None = 0,
) -> MutableMapping[str, float]:
    """Run model inference for one epoch and aggregate detection metrics.

    Args:
        data_loader: Iterable yielding ``(images, targets)`` mini-batches.
        model: Detection model used for inference.
        device: Device where inference should execute.
        num_classes: Total number of classes in evaluation.
        iou_threshold: IoU threshold used for matching in core metrics.
        background_class: Optional class id to exclude from evaluation.

    Returns:
        Mapping containing aggregate precision, recall, F1, and mAP metrics.
    """
    all_outputs: list[Mapping[str, object]] = []
    all_targets: list[Mapping[str, object]] = []

    model.eval()
    with torch.no_grad():
        for images, targets in data_loader:
            images = [image.to(device) for image in images]
            batch_targets = [
                {key: value.to(device) if isinstance(value, torch.Tensor) else value for key, value in item.items()}
                for item in targets
            ]

            batch_outputs = model(images)
            all_outputs.extend(batch_outputs)
            all_targets.extend(batch_targets)

    precision, recall, f1_score = compute_metrics(
        outputs=all_outputs,
        targets=all_targets,
        num_classes=num_classes,
        iou_threshold=iou_threshold,
        background_class=background_class,
    )
    map50, map50_95, _ = compute_map_metrics(
        outputs=all_outputs,
        targets=all_targets,
        num_classes=num_classes,
        background_class=background_class,
    )

    return {
        "precision": float(precision.mean().item()) if precision.numel() else 0.0,
        "recall": float(recall.mean().item()) if recall.numel() else 0.0,
        "f1_score": float(f1_score.mean().item()) if f1_score.numel() else 0.0,
        "mAP@50": float(map50),
        "mAP@50-95": float(map50_95),
    }
