import torch
import torch.nn as nn
import torch.nn.functional as F


class AtrousConvBNReLU(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size, dilation=1, padding=0):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size,
                      padding=padding, dilation=dilation, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )
    def forward(self, x):
        return self.block(x)


class ASPP(nn.Module):
    # 확장 공간 피라미드 풀링 (DeepLabV3).
    # in_ch: 백본 출력 채널 (예: ResNet-101의 2048)
    # out_ch: ASPP 출력 채널 (논문의 256)
    # rates: 병렬 3×3 분기의 확장 비율 목록

    def __init__(self, in_ch=2048, out_ch=256, rates=(6, 12, 18)):
        super().__init__()
        # 분기 1: 1×1 합성곱 (r=1, 패딩 불필요)
        self.b1 = AtrousConvBNReLU(in_ch, out_ch, 1)
        # 분기 2,3,4: 3×3 확장 합성곱
        self.branches = nn.ModuleList([
            AtrousConvBNReLU(in_ch, out_ch, 3, dilation=r, padding=r)
            for r in rates
        ])
        # 분기 5: 전역 평균 풀링
        self.global_avg = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_ch, out_ch, 1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )
        # 융합 1×1 합성곱
        num_branches = 1 + len(rates) + 1        # b1 + 확장 + 전역
        self.project = nn.Sequential(
            nn.Conv2d(num_branches * out_ch, out_ch, 1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
        )

    def forward(self, x):
        H, W = x.shape[-2:]
        # 모든 분기 출력 수집
        out = [self.b1(x)]
        for branch in self.branches:
            out.append(branch(x))
        # 전역 분기: 풀 → 합성곱 → 공간 크기로 업샘플링
        g = self.global_avg(x)
        g = F.interpolate(g, size=(H, W), mode='bilinear', align_corners=False)
        out.append(g)
        # 연결 및 투영
        out = torch.cat(out, dim=1)              # B × (5×256) × H × W
        return self.project(out)                  # B × 256 × H × W


# 빠른 테스트
if __name__ == "__main__":
    aspp = ASPP(in_ch=2048, out_ch=256, rates=(6, 12, 18))
    x = torch.randn(2, 2048, 65, 65)            # 전형적인 ResNet-101 출력
    out = aspp(x)
    print(out.shape)                             # torch.Size([2, 256, 65, 65])
    total = sum(p.numel() for p in aspp.parameters())
    print(f"ASPP 매개변수: {total:,}")             # ~16M