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


# Dropout layer that works even in the evaluation mode
class DropoutAlways(nn.Dropout2d):
    def forward(self, x):
        return F.dropout2d(x, self.p, training=True)

class DownBlock(nn.Module):
    def __init__(self, in_channels, out_channels, normalize=True):
        super().__init__()

        self.block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 4, 2, 1, padding_mode='reflect', bias=False if normalize else True),
            nn.InstanceNorm2d(out_channels, affine=True) if normalize else nn.Identity(),
            # Note that nn.Identity() is just a placeholder layer that returns its input.
            nn.LeakyReLU(0.2),
        )

    def forward(self, x):
        return self.block(x)


class UpBlock(nn.Module):
    def __init__(self, in_channels, out_channels, normalize=True, dropout=False, activation='relu'):
        super().__init__()

        self.block = nn.Sequential(
            nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1, bias=False if normalize else True),
            nn.InstanceNorm2d(out_channels, affine=True) if normalize else nn.Identity(),
            DropoutAlways() if dropout else nn.Identity(),
            nn.ReLU() if activation == 'relu' else nn.Tanh(),
        )

    def forward(self, x):
        return self.block(x)


class Generator(nn.Module):
    def __init__(self):
        super().__init__()

        # Encoder
        self.encoder1 = DownBlock(1, 64, normalize=False)  # 256x256 -> 128x128
        self.encoder2 = DownBlock(64, 128)  # 128x128 -> 64x64
        self.encoder3 = DownBlock(128, 256)  # 64x64 -> 32x32
        self.encoder4 = DownBlock(256, 512)  # 32x32 -> 16x16
        self.encoder5 = DownBlock(512, 512)  # 16x16 -> 8x8
        self.encoder6 = DownBlock(512, 512)  # 8x8 -> 4x4
        self.encoder7 = DownBlock(512, 512)  # 4x4 -> 2x2
        self.encoder8 = DownBlock(512, 512, normalize=False)  # 2x2 -> 1x1

        # Decoder
        self.decoder1 = UpBlock(512, 512, dropout=True)  # 1x1 -> 2x2
        self.decoder2 = UpBlock(512 * 2, 512, dropout=True)  # 2x2 -> 4x4
        self.decoder3 = UpBlock(512 * 2, 512, dropout=True)  # 4x4 -> 8x8
        self.decoder4 = UpBlock(512 * 2, 512)  # 8x8 -> 16x16
        self.decoder5 = UpBlock(512 * 2, 256)  # 16x16 -> 32x32
        self.decoder6 = UpBlock(256 * 2, 128)  # 32x32 -> 64x64
        self.decoder7 = UpBlock(128 * 2, 64)  # 64x64 -> 128x128
        self.decoder8 = UpBlock(64 * 2, 2, normalize=False, activation='tanh')  # 128x128 -> 256x256

    def forward(self, x):
        # Encoder
        ch256_down = x
        ch128_down = self.encoder1(ch256_down)
        ch64_down = self.encoder2(ch128_down)
        ch32_down = self.encoder3(ch64_down)
        ch16_down = self.encoder4(ch32_down)
        ch8_down = self.encoder5(ch16_down)
        ch4_down = self.encoder6(ch8_down)
        ch2_down = self.encoder7(ch4_down)
        ch1 = self.encoder8(ch2_down)

        # Decoder
        ch2_up = self.decoder1(ch1)
        ch2 = torch.cat([ch2_up, ch2_down], dim=1)
        ch4_up = self.decoder2(ch2)
        ch4 = torch.cat([ch4_up, ch4_down], dim=1)
        ch8_up = self.decoder3(ch4)
        ch8 = torch.cat([ch8_up, ch8_down], dim=1)
        ch16_up = self.decoder4(ch8)
        ch16 = torch.cat([ch16_up, ch16_down], dim=1)
        ch32_up = self.decoder5(ch16)
        ch32 = torch.cat([ch32_up, ch32_down], dim=1)
        ch64_up = self.decoder6(ch32)
        ch64 = torch.cat([ch64_up, ch64_down], dim=1)
        ch128_up = self.decoder7(ch64)
        ch128 = torch.cat([ch128_up, ch128_down], dim=1)
        ch256_up = self.decoder8(ch128)

        return ch256_up