"""缺陷分类 + 定位主模型。"""

from __future__ import annotations

from collections.abc import Mapping
from typing import Any

import torch
from torch import nn

from models.heads.center_head import CenterHead
from models.heads.classification import ClassificationHead
from models.registry import BACKBONES, MODELS, NECKS
from utils.labels import LABEL_MODES, MULTILABEL, check_label_mode

__all__ = ["DefectNet", "classification_outputs", "LABEL_MODES"]


def classification_outputs(num_classes: int, label_mode: str) -> int:
    """分类头的输出维度；multiclass 需要额外一维表示 "无缺陷"。"""
    check_label_mode(label_mode)
    return num_classes if label_mode == MULTILABEL else num_classes + 1


@MODELS.register("defect_net")
class DefectNet(nn.Module):
    """backbone -> neck -> {分类头, 中心点定位头}。

    Args:
        num_classes: 缺陷类别数（不含背景）。
        backbone: backbone 注册表的配置，如 `{"name": "tiny_conv", ...}`。
        neck: neck 配置，`None` 表示直接在骨架特征上出图。
        head: CenterHead 的额外参数。
        with_classification: 是否输出图像级分类 logits。
        label_mode: 图像级标签形式，由 yaml 配置决定。
            "multilabel" —— 多类：一张图可同时含多种缺陷，逐类 sigmoid + BCE；
            "multiclass" —— 单类：一张图只归一个类别（含 "无缺陷"），softmax + CE。
            仅影响分类头与分类损失；定位头始终按 num_classes 个类别输出 heatmap。
    """

    def __init__(
        self,
        num_classes: int,
        backbone: Mapping[str, Any] | None = None,
        neck: Mapping[str, Any] | None = None,
        head: Mapping[str, Any] | None = None,
        with_classification: bool = True,
        classification: Mapping[str, Any] | None = None,
        label_mode: str = "multilabel",
    ) -> None:
        super().__init__()
        self.num_classes = num_classes
        self.label_mode = check_label_mode(label_mode)

        self.backbone = BACKBONES.build(dict(backbone or {"name": "tiny_conv"}))
        stride = self.backbone.out_stride
        channels = self.backbone.out_channels

        if neck:
            self.neck: nn.Module = NECKS.build(dict(neck), in_channels=channels)
            channels = self.neck.out_channels
            stride = stride // self.neck.scale_factor
        else:
            self.neck = nn.Identity()

        self.out_stride = stride
        self.center_head = CenterHead(channels, num_classes, **dict(head or {}))
        self.cls_head = (
            ClassificationHead(
                channels,
                classification_outputs(num_classes, label_mode),
                **dict(classification or {}),
            )
            if with_classification
            else None
        )

    def forward(self, images: torch.Tensor) -> dict[str, torch.Tensor]:
        feat = self.neck(self.backbone(images))
        outputs = self.center_head(feat)
        if self.cls_head is not None:
            outputs["logits"] = self.cls_head(feat)
        return outputs