"""缺陷分类模型：卷积骨架 + 分类头。

与 `DefectNet` 共用同一套 backbone 注册表，所以自监督预训练出来的骨架
（`runs/*/simclr_backbone.pth`）可以直接拿来初始化——那正是前面聚类阶段的产物。
按项目约定只用 convolution，不含 attention / transformer。
"""

from __future__ import annotations

from collections.abc import Mapping
from pathlib import Path
from typing import Any

import torch
from torch import nn

from models.heads.classification import ClassificationHead
from models.registry import BACKBONES, MODELS
from utils.logging import get_logger

logger = get_logger("models.classifier")


@MODELS.register("classifier")
class DefectClassifier(nn.Module):
    """Args:
    num_classes: 类别数。
    backbone: backbone 注册表的配置。
    classification: 分类头参数（`hidden`/`dropout` 等）。
    init_from: 自监督骨架权重路径；只加载 backbone，分类头仍随机初始化。
    frozen_backbone: 冻结骨架只训分类头（线性探针，用来单独衡量自监督特征的好坏）。
    """

    def __init__(
        self,
        num_classes: int,
        backbone: Mapping[str, Any] | None = None,
        classification: Mapping[str, Any] | None = None,
        init_from: str | Path | None = None,
        frozen_backbone: bool = False,
    ) -> None:
        super().__init__()
        self.num_classes = num_classes
        self.backbone = BACKBONES.build(dict(backbone or {"name": "tiny_conv"}))
        self.head = ClassificationHead(
            self.backbone.out_channels, num_classes, **dict(classification or {})
        )

        if init_from:
            self.load_backbone(init_from)
        if frozen_backbone:
            for param in self.backbone.parameters():
                param.requires_grad_(False)
            logger.info("骨架已冻结，只训练分类头（线性探针）")

    def load_backbone(self, path: str | Path) -> None:
        """加载自监督预训练的骨架权重。

        只取 backbone 那部分：simclr_backbone.pth 里存的就是骨架 state_dict，
        分类头是新的、必须从头学。
        """
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"预训练骨架不存在: {path}")
        state = torch.load(path, map_location="cpu", weights_only=False)
        weights = state.get("model", state.get("backbone", state))
        missing, unexpected = self.backbone.load_state_dict(weights, strict=False)
        logger.info(
            "已从 %s 加载骨架权重（缺失 %d 项，多余 %d 项，已训 %s 轮）",
            path,
            len(missing),
            len(unexpected),
            state.get("epochs", "?"),
        )
        if len(missing) > len(list(self.backbone.state_dict())) // 2:
            raise ValueError(
                f"{path} 与当前骨架结构对不上（缺失 {len(missing)} 项）；"
                "确认 backbone 配置与自监督训练时一致"
            )

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        return self.head(self.backbone(images))