"""自研轻量卷积骨架，面向实时推理。"""

from __future__ import annotations

from collections.abc import Sequence

import torch
from torch import nn

from models.layers import ConvBNAct, DepthwiseSeparableConv, ResidualBlock
from models.registry import BACKBONES

_BLOCKS = {"dw": DepthwiseSeparableConv, "res": ResidualBlock}


@BACKBONES.register("tiny_conv")
class TinyConvNet(nn.Module):
    """4 个下采样 stage，输出 stride=16 的单尺度特征。

    Args:
        in_channels: 输入通道（薄膜灰度图为 1，RGB 为 3）。
        widths: 每个 stage 的输出通道。
        depths: 每个 stage 的 block 数量。
        block: "dw"(深度可分离，更快) 或 "res"(残差，更强)。
    """

    def __init__(
        self,
        in_channels: int = 3,
        widths: Sequence[int] = (32, 64, 128, 256),
        depths: Sequence[int] = (1, 2, 2, 2),
        block: str = "dw",
    ) -> None:
        super().__init__()
        if len(widths) != len(depths):
            raise ValueError("widths 与 depths 长度必须一致")
        if block not in _BLOCKS:
            raise ValueError(f"未知 block 类型: {block}，可选 {sorted(_BLOCKS)}")
        block_cls = _BLOCKS[block]

        self.stem = ConvBNAct(in_channels, widths[0], kernel_size=3, stride=2)

        stages: list[nn.Module] = []
        prev = widths[0]
        for stage_idx, (width, depth) in enumerate(zip(widths, depths, strict=True)):
            layers: list[nn.Module] = []
            for block_idx in range(depth):
                stride = 2 if (block_idx == 0 and stage_idx > 0) else 1
                layers.append(block_cls(prev, width, stride))
                prev = width
            stages.append(nn.Sequential(*layers))
        self.stages = nn.ModuleList(stages)

        self.out_channels = prev
        # stem 下采样 2 倍，其后每个 stage(除第一个)各下采样 2 倍。
        self.out_stride = 2 * 2 ** (len(widths) - 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        for stage in self.stages:
            x = stage(x)
        return x