"""
自己搭建的网络块文件
包括：
1、基本残差网络（分前向激活和非前向激活版本）\
2、瓶颈结构残差块（Bottleneck Block）（分前向激活和非前向激活版本）\
"""

class ResidualBlock(nn.Module):
"""
基础残差块（ResNet风格）
包含跳跃连接（Skip Connection），解决梯度消失问题
包含前向激活（preactivated）,解决100层以上网络训练困难问题
一次添加两个-三个（多一个下采样）卷积层

*#初始化定义\*
def __init__(self, in_channels, out_channels, stride=1,downsample=None, activation=**'relu',preactivated=False):
"""
tip:输入通道必须和输出通道匹配，不然，没法相加 : y = f(x) + x
[f(x)与x的通道必须匹配]
in_channels：输入通道数
out_channels：输出通道数
stride：步长
downsample:
下采样层（尺寸/通道不匹配时使用），需要时传入卷积块，不需要为None,默认None\
activation：激活层，只有relu和leaky_relu\
preactivated:是否前向激活,需要为True,不需要为False（Pre-activation/ResNet-v2风格）*

if not preactivated:
            # ========== 标准 ResNet (ResNet-v1/v1.5) ==========
            """
            顺序：
            Conv(1x1) → BN → ReLU →
            Conv(3x3, stride) → BN → ReLU →  ← stride在这里（ResNet-v1.5推荐）
            Conv(1x1) → BN → Add → ReLU
            """

            self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=1,
                                   stride=1, padding=0, bias=False)  # 降维
            self.bn1 = nn.BatchNorm2d(mid_channels)

            self.conv2 = nn.Conv2d(mid_channels, mid_channels, kernel_size=3,
                                   stride=stride, padding=1, bias=False)  # 特征提取+下采样
            self.bn2 = nn.BatchNorm2d(mid_channels)


            self.conv3 = nn.Conv2d(mid_channels, out_channels, kernel_size=1,
                                   stride=1, padding=0, bias=False)  # 升维
            self.bn3 = nn.BatchNorm2d(out_channels)

        else:
            # ========== 预激活 ResNet-v2 (Pre-activation) ==========
            """
            顺序：
            BN → ReLU → Conv(1x1) →
            BN → ReLU → Conv(3x3, stride) →
            BN → ReLU → Conv(1x1) → Add
            """
            self.bn1 = nn.BatchNorm2d(in_channels)
            self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=1,
                                   stride=1, padding=0, bias=False)

            self.bn2 = nn.BatchNorm2d(mid_channels)
            self.conv2 = nn.Conv2d(mid_channels, mid_channels, kernel_size=3,
                                   stride=stride, padding=1, bias=False)

            self.bn3 = nn.BatchNorm2d(mid_channels)
            self.conv3 = nn.Conv2d(mid_channels, out_channels, kernel_size=1,
                                   stride=1, padding=0, bias=False)

if __name__ == '__main__':

    # ==================== 像 tf.keras.Sequential 一样直接搭 ====================
    # ============================== 模型定义 ===================================
    #使用 OrderedDict（给每层命名，更像TF）

    model = nn.Sequential(OrderedDict([
    ('resBlok', ResidualBottleneckBlock(512,30,512,preactivated=True)),
    ('dowmSample',nn.Conv2d(512, 256, kernel_size=1,
                                   stride=2, padding=0, bias=False))
    ]))

    # ==========================================================================

    # 可以像字典一样访问
    #print(model.block1)  # 访问特定层

    # 直接使用
    #x = torch.randn(2, 3, 224, 224)
    #output = model(x)  # 自动顺序执行，无需forward
    #print(output.shape)  # [2, 10]

    # ================================== 打印模型架构 ============================

    # 1、类似 model.summary() 的表格输出
    #summary(model, input_size=(512, 224, 224), device='cpu')#（输入通道，输入w尺寸，输入h尺寸）

    # ================================== 打印模型架构 ============================
    # 2、torchinfo（推荐，信息更丰富）
        # 输出包含：
        # - 每层的输入/输出形状
        # - 参数量
        # - 乘加运算量（FLOPs估算）
        # - 内存占用
    '''
    sum(model,
        input_size=(2,512, 224, 224),  # (batch, channels, H, W)
        col_names=["input_size", "output_size", "num_params", "kernel_size", "mult_adds"],#显示的列
        col_width=20,#每列显示的宽度
        row_settings=["var_names"],#行的显示设置
        verbose=1)
    '''