# Fucus的作用是将输入的特征图的宽和高信息聚合到通道维度上
# 具体而言是通过将输入的特征图分成4个部分，然后将4个部分的特征图进行concat操作
# 然后将concat后的特征图进行卷积操作
class Focus(nn.Module):
    """Focus width and height information into channel space."""

    def __init__(self, in_channels, out_channels, ksize=1, stride=1, act="silu"):
        super().__init__()
        self.conv = BaseConv(in_channels * 4, out_channels, ksize, stride, act=act)
        # in_channels * 4是因为将输入的特征图分成4个部分，然后将4个部分的特征图进行concat操作
        # 所以输入的通道数是4倍的输出通道数
        # ksize是卷积核的大小，为1，作用是减少参数的数量
        # stride是步长，为1，作用是减少特征图的大小

    def forward(self, x):
        # shape of x (b,c,w,h) -> y(b,4c,w/2,h/2)
        patch_top_left = x[..., ::2, ::2]
        # [..., ::2, ::2]表示取x的第三维和第四维的数据
        # 第三维的数据是从0开始，每隔一个取一个，直到最后一个
        # 第四维的数据是从0开始，每隔一个取一个，直到最后一个
        # 例如：
        # x = torch.arange(16).reshape(1, 1, 4, 4)
        # 表示把arrange(16)的数据变成4x4的矩阵
        # x = tensor([[[[ 0,  1,  2,  3],
        #               [ 4,  5,  6,  7],
        #               [ 8,  9, 10, 11],
        #               [12, 13, 14, 15]]]])
        # patch_top_left = x[..., ::2, ::2]
        # 表示取x的第三维和第四维的数据
        # 第三维的数据是从0开始，每隔一个取一个，直到最后一个
        # 第四维的数据是从0开始，每隔一个取一个，直到最后一个
        # patch_top_left = tensor([[[[ 0,  2],
        #                            [ 8, 10]]]])
        patch_top_right = x[..., ::2, 1::2]
        patch_bot_left = x[..., 1::2, ::2]
        patch_bot_right = x[..., 1::2, 1::2]
        x = torch.cat(
            (
                patch_top_left,
                patch_bot_left,
                patch_top_right,
                patch_bot_right,
            ),
            dim=1, # 在通道维度上进行concat操作
                   # 如果dim=0，表示在batch维度上进行concat操作
                   # 如果dim=2，表示在宽维度上进行concat操作
        )
        return self.conv(x)
# Focus的结构图如下：
# 1. 输入的特征图的大小为HxWxC_in
# 2. 将输入的特征图分成4个部分，然后将4个部分的特征图进行concat操作
# 3. 将2的输出的特征图进行1x1卷积操作，输出的特征图的大小为H/2xW/2xC_out

# Q: 如果输入的特征图的宽和高不是偶数，那么怎么concat呢？
# A: 如果输入的特征图的宽和高不是偶数，那么就将输入的特征图的宽和高减1，然后再进行concat操作