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

class ResnetBlock(nn.Module):
    def __init__(self, dim, padding_type='reflect', norm_layer=nn.BatchNorm2d, use_dropout=False, use_bias=False):
        super(ResnetBlock, self).__init__()
        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)

    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
        conv_block = []
        p = 0
        if padding_type == 'reflect':
            conv_block += [nn.ReflectionPad2d(1)]
        elif padding_type == 'replicate':
            conv_block += [nn.ReplicationPad2d(1)]
        elif padding_type == 'zero':
            p = 1
        else:
            raise NotImplementedError('padding [%s] is not implemented' % padding_type)

        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
                       norm_layer(dim),
                       nn.ReLU(True)]
        if use_dropout:
            conv_block += [nn.Dropout(0.5)]

        p = 0
        if padding_type == 'reflect':
            conv_block += [nn.ReflectionPad2d(1)]
        elif padding_type == 'replicate':
            conv_block += [nn.ReplicationPad2d(1)]
        elif padding_type == 'zero':
            p = 1
        else:
            raise NotImplementedError('padding [%s] is not implemented' % padding_type)

        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
                       norm_layer(dim)]

        return nn.Sequential(*conv_block)

    def forward(self, x):
        out = x + self.conv_block(x)
        return out


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

        # Encoder layers for 84x84 input
        encoder_lis = [
            nn.Conv2d(gen_input_nc, 64, kernel_size=4, stride=2, padding=1, bias=True),  # -> 64x42x42
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=True),          # -> 128x21x21
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=True),         # -> 256x10x10
            nn.BatchNorm2d(256),
            nn.ReLU(),
        ]

        # Bottle-neck with Resnet blocks
        bottle_neck_lis = [
            ResnetBlock(256, norm_layer=nn.BatchNorm2d),
            ResnetBlock(256, norm_layer=nn.BatchNorm2d),
            ResnetBlock(256, norm_layer=nn.BatchNorm2d),
            ResnetBlock(256, norm_layer=nn.BatchNorm2d)
        ]

        # Decoder layers to match 84x84 output
        decoder_lis = [
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, output_padding=0, bias=False),  # -> 128x21x21
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, output_padding=0, bias=False),   # -> 64x42x42
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.ConvTranspose2d(64, image_nc, kernel_size=4, stride=2, padding=1, output_padding=1, bias=False),  # -> image_nc x 84 x 84
            nn.Tanh()
        ]

        self.encoder = nn.Sequential(*encoder_lis)
        self.bottle_neck = nn.Sequential(*bottle_neck_lis)
        self.decoder = nn.Sequential(*decoder_lis)

    def forward(self, x):
        x = self.encoder(x)
        x = self.bottle_neck(x)
        x = self.decoder(x)

        if x.shape[2] != 84 or x.shape[3] != 84:
            x = F.pad(x, (0, 3, 0, 3)) 

        return x



class Discriminator(nn.Module):
    def __init__(self, image_nc):
        super(Discriminator, self).__init__()

        self.first_layer = nn.Sequential(
            nn.Conv2d(image_nc, 16, kernel_size=4, stride=2, padding=1, bias=True),
            nn.LeakyReLU(0.2)
        )

        self.model = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=4, stride=2, padding=1, bias=True),  # -> 32x21x21
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2),
            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1, bias=True),  # -> 64x11x11
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=True),  # -> 128x6x6
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=True),  # -> 256x3x3
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2),
            nn.Conv2d(256, 1, kernel_size=2, stride=1, padding=0, bias=True),  # -> 1x1x1
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.first_layer(x)
        x = self.model(x)
        output = x.squeeze()
        if output.dim() == 0:
            output = output.view(1)
        return output
