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


class AttentionNet(nn.Module):
    """Predicts a soft attention mask for rainy regions."""

    def __init__(self, in_channels: int = 3, base_channels: int = 32, max_channels: int = 128) -> None:
        super().__init__()
        mid_channels = min(base_channels * 2, max_channels)
        bottleneck_channels = min(base_channels * 4, max_channels)

        # Encoder
        self.enc1 = nn.Sequential(
            nn.Conv2d(in_channels, base_channels, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(base_channels, base_channels, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.down1 = nn.Conv2d(base_channels, mid_channels, kernel_size=4, stride=2, padding=1)

        self.enc2 = nn.Sequential(
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(mid_channels, mid_channels, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.down2 = nn.Conv2d(mid_channels, bottleneck_channels, kernel_size=4, stride=2, padding=1)

        self.bottleneck = nn.Sequential(
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(bottleneck_channels, bottleneck_channels, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
        )

        # Decoder with skip connections
        self.up1 = nn.ConvTranspose2d(bottleneck_channels, mid_channels, kernel_size=4, stride=2, padding=1)
        self.dec1 = nn.Sequential(
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(mid_channels * 2, mid_channels, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
        )

        self.up2 = nn.ConvTranspose2d(mid_channels, base_channels, kernel_size=4, stride=2, padding=1)
        self.dec2 = nn.Sequential(
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(base_channels * 2, base_channels, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
        )

        self.out_conv = nn.Conv2d(base_channels, 1, kernel_size=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        enc1 = self.enc1(x)
        enc2 = self.enc2(self.down1(enc1))
        bottleneck = self.bottleneck(self.down2(enc2))

        up1 = self.up1(bottleneck)
        up1 = torch.cat([up1, enc2], dim=1)
        up1 = self.dec1(up1)

        up2 = self.up2(up1)
        up2 = torch.cat([up2, enc1], dim=1)
        up2 = self.dec2(up2)

        return torch.sigmoid(self.out_conv(up2))
