import torch
from torch import nn

class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(ConvBlock, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
        self.batchnorm = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.batchnorm(self.conv(x)))

class DeconvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding):
        super(DeconvBlock, self).__init__()
        self.deconv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding)
        self.batchnorm = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.batchnorm(self.deconv(x)))

class Autoencoder(nn.Module):
    def __init__(self, feature_dim=32):
        super(Autoencoder, self).__init__()
        self.feature_dim = feature_dim

        # エンコーダ
        self.enc1 = ConvBlock(3, 16, 10, 1, 0)
        self.enc2 = ConvBlock(16, 32, 10, 1, 0)
        self.enc3 = ConvBlock(32, 64, 2, 2, 0)
        self.enc4 = ConvBlock(64, 128, 2, 2, 0)
        self.enc5 = ConvBlock(128, 256, 2, 2, 0)

        # デコーダ
        self.dec1 = DeconvBlock(256, 128, 2, 2, 0, 1)
        self.dec2 = DeconvBlock(256, 64, 2, 2, 0, 1) # 128 + 128
        self.dec3 = DeconvBlock(128, 32, 2, 2, 0, 0) # 64 + 64
        self.dec4 = DeconvBlock(64, 16, 10, 1, 0, 0) # 32 + 32
        self.dec5 = DeconvBlock(32, self.feature_dim, 10, 1, 0, 0)
        self.dec6 = nn.Conv2d(self.feature_dim, 32, 1, 1, 0)
        self.dec7 = nn.Conv2d(32, 3, 1, 1, 0)

    def forward(self, x):
        # エンコーダ
        enc1 = self.enc1(x)
        enc2 = self.enc2(enc1)
        enc3 = self.enc3(enc2)
        enc4 = self.enc4(enc3)
        enc5 = self.enc5(enc4)

        # デコーダ
        dec1 = self.dec1(enc5)
        dec2 = self.dec2(torch.cat((dec1, enc4), 1))
        dec3 = self.dec3(torch.cat((dec2, enc3), 1))
        dec4 = self.dec4(torch.cat((dec3, enc2), 1))
        dec5 = self.dec5(torch.cat((dec4, enc1), 1))
        dec6 = self.dec6(dec5)
        dec7 = self.dec7(dec6)

        return dec5, dec7