"""torchvision 卷积骨架封装（resnet / mobilenet 等纯卷积模型）。"""

from __future__ import annotations

import torch
from torch import nn

from models.registry import BACKBONES

# 仅允许纯卷积骨架；ViT/Swin 等不在部署范围内。
_SUPPORTED = {
    "resnet18": 512,
    "resnet34": 512,
    "resnet50": 2048,
    "mobilenet_v3_small": 576,
    "mobilenet_v3_large": 960,
    "shufflenet_v2_x1_0": 1024,
}


@BACKBONES.register("torchvision")
class TorchvisionBackbone(nn.Module):
    """输出 stride=32 的单尺度特征。

    Args:
        arch: `_SUPPORTED` 中的模型名。
        pretrained: 是否加载 ImageNet 权重（需要联网或本地缓存）。
        in_channels: 输入通道，非 3 时会把第一层卷积权重按通道求均值复制。
        frozen: 是否冻结骨架参数。
    """

    def __init__(
        self,
        arch: str = "resnet18",
        pretrained: bool = True,
        in_channels: int = 3,
        frozen: bool = False,
    ) -> None:
        super().__init__()
        if arch not in _SUPPORTED:
            raise ValueError(f"不支持的骨架 {arch}，可选 {sorted(_SUPPORTED)}")

        import torchvision

        weights = "DEFAULT" if pretrained else None
        net = getattr(torchvision.models, arch)(weights=weights)

        if arch.startswith("resnet"):
            self.features = nn.Sequential(*list(net.children())[:-2])
        else:
            self.features = net.features

        if in_channels != 3:
            self._patch_first_conv(in_channels)

        self.out_channels = _SUPPORTED[arch]
        self.out_stride = 32

        if frozen:
            for param in self.parameters():
                param.requires_grad_(False)

    def _patch_first_conv(self, in_channels: int) -> None:
        for module in self.features.modules():
            if isinstance(module, nn.Conv2d):
                weight = module.weight.data.mean(dim=1, keepdim=True).repeat(1, in_channels, 1, 1)
                module.in_channels = in_channels
                module.weight = nn.Parameter(weight * (3.0 / in_channels))
                return
        raise RuntimeError("骨架中未找到 Conv2d，无法调整输入通道")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.features(x)