import torch
import torch.nn as nn

class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.InstanceNorm2d(channels),
            nn.ReLU(),
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.InstanceNorm2d(channels)
        )

    def forward(self, x):
        return x + self.conv(x)
    
class UpsampleConvLayer(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride):
        super().__init__()
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        # Reflection padding reduces "border" artifacts
        self.reflection_pad = nn.ReflectionPad2d(kernel_size // 2)
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=1)

    def forward(self, x):
        x = self.upsample(x)
        x = self.reflection_pad(x)
        return self.conv(x)

class StyleNet(nn.Module):
    def __init__(self):
        super(StyleNet, self).__init__()
        # Encoder
        self.encoder = nn.Sequential(
            nn.ReflectionPad2d(4),           # Pad first...
            nn.Conv2d(3, 32, 9, padding=0),  # ...then Conv (padding=0)
            nn.InstanceNorm2d(32), nn.ReLU(),
            
            nn.ReflectionPad2d(1),
            nn.Conv2d(32, 64, 3, stride=2, padding=0),
            nn.InstanceNorm2d(64), nn.ReLU(),
            
            nn.ReflectionPad2d(1),
            nn.Conv2d(64, 128, 3, stride=2, padding=0),
            nn.InstanceNorm2d(128), nn.ReLU()
        )
        # Bottleneck
        self.bottleneck = nn.Sequential(*[ResidualBlock(128) for _ in range(3)])
        # Decoder
        self.decoder = nn.Sequential(
            UpsampleConvLayer(128, 64, kernel_size=3, stride=1),
            nn.InstanceNorm2d(64), nn.ReLU(),
            UpsampleConvLayer(64, 32, kernel_size=3, stride=1),
            nn.InstanceNorm2d(32), nn.ReLU(),
            nn.Conv2d(32, 3, 9, padding=4),
            nn.Sigmoid()
        )

    def forward(self, x):
        features = self.encoder(x)
        features = self.bottleneck(features)
        return self.decoder(features)