"""把框标注编码成 CenterHead 需要的监督信号。"""

from __future__ import annotations

import numpy as np
import torch


def gaussian_radius(height: float, width: float, min_overlap: float = 0.7) -> float:
    """CenterNet 的半径公式：保证半径内的框与 GT 的 IoU >= min_overlap。"""
    a1 = 1.0
    b1 = height + width
    c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
    r1 = (b1 - np.sqrt(max(b1**2 - 4 * a1 * c1, 0.0))) / (2 * a1)

    a2 = 4.0
    b2 = 2 * (height + width)
    c2 = (1 - min_overlap) * width * height
    r2 = (b2 - np.sqrt(max(b2**2 - 4 * a2 * c2, 0.0))) / (2 * a2)

    a3 = 4 * min_overlap
    b3 = -2 * min_overlap * (height + width)
    c3 = (min_overlap - 1) * width * height
    r3 = (-b3 + np.sqrt(max(b3**2 - 4 * a3 * c3, 0.0))) / (2 * a3)

    return max(0.0, min(r1, r2, r3))


def draw_gaussian(heatmap: np.ndarray, center: tuple[int, int], radius: int) -> None:
    """在 heatmap 上叠加二维高斯（取 max，避免相邻目标互相削弱）。"""
    diameter = 2 * radius + 1
    sigma = diameter / 6.0
    grid = np.arange(-radius, radius + 1, dtype=np.float32)
    gaussian = np.exp(-(grid[None, :] ** 2 + grid[:, None] ** 2) / (2 * sigma**2))

    cx, cy = center
    height, width = heatmap.shape
    left, right = min(cx, radius), min(width - cx, radius + 1)
    top, bottom = min(cy, radius), min(height - cy, radius + 1)
    if right <= -left or bottom <= -top:
        return

    masked_heatmap = heatmap[cy - top : cy + bottom, cx - left : cx + right]
    masked_gaussian = gaussian[radius - top : radius + bottom, radius - left : radius + right]
    np.maximum(masked_heatmap, masked_gaussian, out=masked_heatmap)


class CenterTargetEncoder:
    """生成 heatmap / wh / offset / mask 四个监督张量。

    Args:
        num_classes: 类别数。
        stride: 模型输出相对输入的下采样倍数。
        max_objects: 单图最多编码的目标数。
        min_overlap: 高斯半径的 IoU 阈值。
    """

    def __init__(
        self,
        num_classes: int,
        stride: int,
        max_objects: int = 64,
        min_overlap: float = 0.7,
    ) -> None:
        self.num_classes = num_classes
        self.stride = stride
        self.max_objects = max_objects
        self.min_overlap = min_overlap

    def __call__(
        self,
        boxes: np.ndarray,
        labels: np.ndarray,
        image_size: tuple[int, int],
    ) -> dict[str, torch.Tensor]:
        img_h, img_w = image_size
        out_h, out_w = img_h // self.stride, img_w // self.stride

        heatmap = np.zeros((self.num_classes, out_h, out_w), dtype=np.float32)
        wh = np.zeros((2, out_h, out_w), dtype=np.float32)
        offset = np.zeros((2, out_h, out_w), dtype=np.float32)
        mask = np.zeros((out_h, out_w), dtype=np.float32)

        count = 0
        for box, label in zip(np.asarray(boxes, dtype=np.float32), np.asarray(labels), strict=True):
            if count >= self.max_objects:
                break
            x1, y1, x2, y2 = box / self.stride
            box_w, box_h = x2 - x1, y2 - y1
            if box_w <= 0 or box_h <= 0:
                continue

            cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
            cx_int, cy_int = int(cx), int(cy)
            if not (0 <= cx_int < out_w and 0 <= cy_int < out_h):
                continue

            radius = max(1, int(gaussian_radius(box_h, box_w, self.min_overlap)))
            draw_gaussian(heatmap[int(label)], (cx_int, cy_int), radius)

            wh[:, cy_int, cx_int] = (box_w, box_h)
            offset[:, cy_int, cx_int] = (cx - cx_int, cy - cy_int)
            mask[cy_int, cx_int] = 1.0
            count += 1

        return {
            "heatmap": torch.from_numpy(heatmap),
            "wh": torch.from_numpy(wh),
            "offset": torch.from_numpy(offset),
            "mask": torch.from_numpy(mask),
        }