"""推理后处理：从 heatmap 解码出缺陷框。"""

from __future__ import annotations

import torch
from torch.nn import functional as F


def heatmap_nms(heatmap: torch.Tensor, kernel: int = 3) -> torch.Tensor:
    """用 max-pool 做峰值抑制，替代 NMS（无循环，可导出）。"""
    pooled = F.max_pool2d(heatmap, kernel, stride=1, padding=kernel // 2)
    return heatmap * (pooled == heatmap).float()


@torch.no_grad()
def decode_heatmap(
    heatmap: torch.Tensor,
    wh: torch.Tensor,
    offset: torch.Tensor,
    stride: int,
    topk: int = 50,
    score_threshold: float = 0.3,
) -> list[dict[str, torch.Tensor]]:
    """解码为原图坐标下的检测结果。

    Args:
        heatmap: (B, C, H, W) logits。
        wh / offset: (B, 2, H, W)，单位为特征图像素。
        stride: 特征图相对原图的下采样倍数。

    Returns:
        每张图一个 dict：`boxes` (N, 4) xyxy、`scores` (N,)、`labels` (N,)。
    """
    scores_map = heatmap_nms(heatmap.sigmoid())
    batch, _num_classes, height, width = scores_map.shape

    flat = scores_map.view(batch, -1)
    topk = min(topk, flat.shape[1])
    topk_scores, topk_idx = flat.topk(topk, dim=1)

    labels = torch.div(topk_idx, height * width, rounding_mode="floor")
    pixel_idx = topk_idx % (height * width)
    ys = torch.div(pixel_idx, width, rounding_mode="floor").float()
    xs = (pixel_idx % width).float()

    wh_flat = wh.view(batch, 2, -1)
    offset_flat = offset.view(batch, 2, -1)
    gather_idx = pixel_idx.unsqueeze(1).expand(-1, 2, -1)
    pred_wh = wh_flat.gather(2, gather_idx)
    pred_offset = offset_flat.gather(2, gather_idx)

    cx = (xs + pred_offset[:, 0]) * stride
    cy = (ys + pred_offset[:, 1]) * stride
    half_w = pred_wh[:, 0].clamp(min=0) * stride / 2
    half_h = pred_wh[:, 1].clamp(min=0) * stride / 2
    boxes = torch.stack([cx - half_w, cy - half_h, cx + half_w, cy + half_h], dim=-1)

    results = []
    for i in range(batch):
        keep = topk_scores[i] >= score_threshold
        results.append(
            {
                "boxes": boxes[i][keep],
                "scores": topk_scores[i][keep],
                "labels": labels[i][keep],
            }
        )
    return results