# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
# Copyright (c) 2024 Baidu. All Rights Reserved.
# ------------------------------------------------------------------------
# Conditional DETR
# Copyright (c) 2021 Microsoft. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Copied from DETR (https://github.com/facebookresearch/detr)
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
# ------------------------------------------------------------------------

"""Greedy matching and accumulation functions for evaluation metrics."""

from typing import Any, List

import numpy as np
import torch
import torch.nn.functional as F
from torchvision.ops import box_iou

from rfdetr.utilities import all_gather


def _compute_mask_iou(pred_masks: torch.Tensor, gt_masks: torch.Tensor) -> torch.Tensor:
    """Compute pairwise boolean-mask IoU between N predictions and M ground truths.

    Args:
        pred_masks: Boolean mask tensor of shape [N, H, W].
        gt_masks: Boolean mask tensor of shape [M, H, W].

    Returns:
        IoU tensor of shape [N, M].
    """
    n = pred_masks.shape[0]
    m = gt_masks.shape[0]
    if pred_masks.shape[-2:] != gt_masks.shape[-2:]:
        h, w = pred_masks.shape[-2:]
        gt_masks = F.interpolate(gt_masks.float().unsqueeze(1), size=(h, w), mode="nearest").squeeze(1)
    pred_flat = pred_masks.bool().view(n, -1).float()  # [N, HW]
    gt_flat = gt_masks.bool().view(m, -1).float()  # [M, HW]
    inter = torch.mm(pred_flat, gt_flat.t())  # [N, M]
    pred_area = pred_flat.sum(dim=1, keepdim=True)  # [N, 1]
    gt_area = gt_flat.sum(dim=1, keepdim=True)  # [M, 1]
    union = pred_area + gt_area.t() - inter  # [N, M]
    return torch.where(union > 0, inter / union, torch.zeros_like(inter))


def _match_single_class(
    pred_scores: torch.Tensor,
    pred_items: torch.Tensor,
    gt_items: torch.Tensor,
    gt_crowd: torch.Tensor,
    iou_threshold: float,
    iou_type: str,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, int]:
    """Greedy highest-score-first matching for one class in one image.

    Implements the COCO matching algorithm: each GT is matched at most once;
    detections are processed in descending score order; detections matched to
    crowd GTs are marked as ignored rather than false positives.

    Args:
        pred_scores: Float tensor of shape [N] with detection confidences.
        pred_items: Predictions — boxes [N, 4] in xyxy coords or masks [N, H, W].
        gt_items: Ground truths — boxes [M, 4] in xyxy coords or masks [M, H, W].
        gt_crowd: Bool tensor of shape [M], True for crowd instances.
        iou_threshold: Minimum IoU to count as a positive match.
        iou_type: ``"bbox"`` for box IoU or ``"segm"`` for mask IoU.

    Returns:
        Tuple ``(scores_np, matches_np, ignore_np, total_gt)`` where:
            - scores_np: float32 array [N] ordered by descending score.
            - matches_np: int array [N], 1 = TP, 0 = FP.
            - ignore_np: bool array [N], True if matched to a crowd GT.
            - total_gt: number of non-crowd GT instances.
    """
    n = pred_scores.shape[0]
    m = gt_items.shape[0]

    sort_idx = torch.argsort(pred_scores, descending=True)
    pred_scores_sorted = pred_scores[sort_idx]
    pred_sorted = pred_items[sort_idx]

    if iou_type == "bbox":
        iou_matrix = box_iou(pred_sorted, gt_items)  # [N, M]
    else:
        iou_matrix = _compute_mask_iou(pred_sorted, gt_items)  # [N, M]

    device = pred_scores.device
    gt_matched = torch.zeros(m, dtype=torch.bool, device=device)
    pred_match = torch.zeros(n, dtype=torch.long, device=device)
    pred_ignore = torch.zeros(n, dtype=torch.bool, device=device)

    for i in range(n):
        ious = iou_matrix[i]  # [M]

        # Try to match to a non-crowd GT (each non-crowd GT matched at most once).
        nc_ious = ious.clone()
        nc_ious[gt_crowd] = -1.0
        nc_ious[gt_matched & ~gt_crowd] = -1.0  # already claimed

        best_nc_iou, best_nc_idx = nc_ious.max(dim=0)
        if best_nc_iou >= iou_threshold:
            pred_match[i] = 1
            gt_matched[best_nc_idx] = True
        else:
            # A detection matched to a crowd GT is ignored (not a false positive).
            if gt_crowd.any():
                crowd_ious = ious.clone()
                crowd_ious[~gt_crowd] = -1.0
                if crowd_ious.max() >= iou_threshold:
                    pred_ignore[i] = True
            # else: false positive — pred_match stays 0

    total_gt = int((~gt_crowd).sum().item())
    return (
        pred_scores_sorted.float().cpu().numpy().astype(np.float32),
        pred_match.cpu().numpy(),
        pred_ignore.cpu().numpy().astype(bool),
        total_gt,
    )


def build_matching_data(
    preds_list: list[dict[str, torch.Tensor]],
    targets_list: list[dict[str, torch.Tensor]],
    iou_threshold: float = 0.5,
    iou_type: str = "bbox",
) -> dict[int, dict[str, Any]]:
    """Build compact per-class matching data from a batch of predictions and targets.

    Implements greedy highest-score-first matching compatible with the COCO
    algorithm. The returned dict can be passed directly to
    ``merge_matching_data()`` and ultimately consumed by
    ``sweep_confidence_thresholds()`` after conversion to list form.

    Args:
        preds_list: Per-image predictions. Each dict must contain:

            - ``boxes``: float Tensor [N, 4] in absolute xyxy coordinates.
            - ``scores``: float Tensor [N].
            - ``labels``: int64 Tensor [N].
            - ``masks`` *(optional)*: bool Tensor [N, H, W] for segmentation.

        targets_list: Per-image ground truths. Each dict must contain:

            - ``boxes``: float Tensor [M, 4] in absolute xyxy coordinates.
            - ``labels``: int64 Tensor [M].
            - ``masks`` *(optional)*: bool Tensor [M, H, W] for segmentation.
            - ``iscrowd`` *(optional)*: int64 Tensor [M], 1 for crowd instances.

        iou_threshold: IoU threshold for positive matching. Defaults to 0.5.
        iou_type: ``"bbox"`` for bounding-box IoU; ``"segm"`` for boolean-mask
            IoU. Defaults to ``"bbox"``.

    Returns:
        Dict mapping ``class_id`` (int) to a compact matching dict with keys:

            - ``"scores"``: float32 ndarray of detection scores.
            - ``"matches"``: int ndarray (1 = TP, 0 = FP).
            - ``"ignore"``: bool ndarray (True if matched to a crowd GT).
            - ``"total_gt"``: int, count of non-crowd GT instances.
    """
    acc: dict[int, dict[str, list | int]] = {}

    for preds, targets in zip(preds_list, targets_list):
        pred_boxes = preds["boxes"]  # [N, 4]
        pred_scores = preds["scores"]  # [N]
        pred_labels = preds["labels"]  # [N]
        pred_masks = preds.get("masks")  # [N, H, W] | None

        gt_boxes = targets["boxes"]  # [M, 4]
        gt_labels = targets["labels"]  # [M]
        gt_masks = targets.get("masks")  # [M, H, W] | None
        raw_crowd = targets.get(
            "iscrowd",
            torch.zeros(len(gt_labels), dtype=torch.long, device=gt_labels.device),
        )
        gt_crowd = raw_crowd.bool()

        all_class_ids: set[int] = set(gt_labels.tolist()) | set(pred_labels.tolist())

        for class_id in all_class_ids:
            pred_mask_c = pred_labels == class_id
            gt_mask_c = gt_labels == class_id

            p_scores = pred_scores[pred_mask_c]
            gt_crowd_c = gt_crowd[gt_mask_c]
            n_pred = int(pred_mask_c.sum().item())
            n_gt = int(gt_mask_c.sum().item())

            entry = acc.setdefault(
                class_id,
                {"scores": [], "matches": [], "ignore": [], "total_gt": 0},
            )

            if n_pred == 0:
                entry["total_gt"] += int((~gt_crowd_c).sum().item())
                continue

            if n_gt == 0:
                # TODO: support bfloat16 natively once numpy adds bf16 dtype
                sc = p_scores.float().cpu().numpy()
                order = np.argsort(-sc)
                entry["scores"].extend(sc[order].tolist())
                entry["matches"].extend([0] * n_pred)
                entry["ignore"].extend([False] * n_pred)
                continue

            if iou_type == "bbox":
                p_items: torch.Tensor = pred_boxes[pred_mask_c]  # [n_pred, 4]
                gt_items: torch.Tensor = gt_boxes[gt_mask_c]  # [n_gt, 4]
            else:
                if pred_masks is None or gt_masks is None:
                    raise ValueError("iou_type='segm' requires 'masks' in both preds and targets")
                p_items = pred_masks[pred_mask_c]  # [n_pred, H, W]
                gt_items = gt_masks[gt_mask_c]  # [n_gt, H, W]

            scores_np, matches_np, ignore_np, total_gt = _match_single_class(
                p_scores, p_items, gt_items, gt_crowd_c, iou_threshold, iou_type
            )

            entry["scores"].extend(scores_np.tolist())
            entry["matches"].extend(matches_np.tolist())
            entry["ignore"].extend(ignore_np.tolist())
            entry["total_gt"] += total_gt

    return {
        class_id: {
            "scores": np.array(data["scores"], dtype=np.float32),
            "matches": np.array(data["matches"], dtype=np.int64),
            "ignore": np.array(data["ignore"], dtype=bool),
            "total_gt": data["total_gt"],
        }
        for class_id, data in acc.items()
    }


def init_matching_accumulator() -> dict[int, dict[str, Any]]:
    """Return an empty matching accumulator compatible with ``merge_matching_data()``.

    Returns:
        Empty dict to be passed as the first argument to ``merge_matching_data()``.
    """
    return {}


def merge_matching_data(
    accumulator: dict[int, dict[str, Any]],
    new_data: dict[int, dict[str, Any]],
) -> dict[int, dict[str, Any]]:
    """Merge *new_data* into *accumulator* in place.

    Both arguments share the dict schema produced by ``build_matching_data()``:
    each class-keyed sub-dict contains ``"scores"`` (float32 ndarray),
    ``"matches"`` (int64 ndarray), ``"ignore"`` (bool ndarray), and
    ``"total_gt"`` (int).

    Args:
        accumulator: Running accumulator, modified in place.
        new_data: Batch-level matching data to merge in.

    Returns:
        The modified *accumulator* (same object, for method chaining).
    """
    for class_id, data in new_data.items():
        if class_id not in accumulator:
            accumulator[class_id] = {
                "scores": data["scores"].copy(),
                "matches": data["matches"].copy(),
                "ignore": data["ignore"].copy(),
                "total_gt": data["total_gt"],
            }
        else:
            entry = accumulator[class_id]
            entry["scores"] = np.concatenate([entry["scores"], data["scores"]])
            entry["matches"] = np.concatenate([entry["matches"], data["matches"]])
            entry["ignore"] = np.concatenate([entry["ignore"], data["ignore"]])
            entry["total_gt"] += data["total_gt"]
    return accumulator


def distributed_merge_matching_data(
    local_data: dict[int, dict[str, Any]],
) -> dict[int, dict[str, Any]]:
    """Gather per-rank matching data from all DDP ranks and merge into one dict.

    Uses ``rfdetr.utilities.all_gather`` (pickle-based) so the data need not be a tensor.
    In single-process (non-distributed) mode, returns a merged copy of *local_data*
    unchanged.

    Args:
        local_data: Per-rank accumulator produced by ``merge_matching_data()``.

    Returns:
        Merged accumulator containing contributions from all ranks.
    """
    gathered: List[dict[int, dict[str, Any]]] = all_gather(local_data)
    merged: dict[int, dict[str, Any]] = {}
    for rank_data in gathered:
        merge_matching_data(merged, rank_data)
    return merged
