"""Anchor-free 中心点定位头（CenterNet 风格，纯卷积）。

输出:
    heatmap: (B, num_classes, H, W)  各类别缺陷中心的置信度（logits）
    wh:      (B, 2, H, W)            中心点处的框宽高（特征图尺度）
    offset:  (B, 2, H, W)            下采样量化误差补偿
"""

from __future__ import annotations

import math

import torch
from torch import nn

from models.layers import ConvBNAct
from models.registry import HEADS


def _branch(in_channels: int, hidden: int, out_channels: int) -> nn.Sequential:
    return nn.Sequential(
        ConvBNAct(in_channels, hidden, kernel_size=3),
        nn.Conv2d(hidden, out_channels, kernel_size=1),
    )


@HEADS.register("center")
class CenterHead(nn.Module):
    def __init__(
        self,
        in_channels: int,
        num_classes: int,
        hidden_channels: int = 64,
        prior_prob: float = 0.01,
    ) -> None:
        super().__init__()
        self.num_classes = num_classes
        self.heatmap = _branch(in_channels, hidden_channels, num_classes)
        self.wh = _branch(in_channels, hidden_channels, 2)
        self.offset = _branch(in_channels, hidden_channels, 2)

        # heatmap 正样本极稀疏，用先验概率初始化最后一层 bias，稳定早期训练。
        bias = -math.log((1 - prior_prob) / prior_prob)
        nn.init.constant_(self.heatmap[-1].bias, bias)

    def forward(self, feat: torch.Tensor) -> dict[str, torch.Tensor]:
        return {
            "heatmap": self.heatmap(feat),
            "wh": self.wh(feat),
            "offset": self.offset(feat),
        }