import torch.nn as nn
class ResidualBlock(nn.Module):
    def __init__(self,input_channels ) -> None:
        super(ResidualBlock,self).__init__()
        self.conv1 = nn.Conv2d(input_channels,input_channels,3,1,padding=1,padding_mode='reflect')
        self.conv2 = nn.Conv2d(input_channels,input_channels,3,1,padding=1,padding_mode='reflect')
        self.instanceNorm = nn.InstanceNorm2d(input_channels)
        self.activation = nn.ReLU()

    def forward(self,x):
        original = x.copy()
        x = self.conv1(x)
        x = self.instanceNorm(x)
        x = self.activation(x)
        x = self.conv2(x)
        x = self.instanceNorm(x)
        return original + x
    


class ContractingBlock(nn.Module):
    def __init__(self, input_channels, use_bn=True,kernel_size=3,activation='relu') -> None:
        super(ContractingBlock,self).__init__()
        self.conv1 = nn.Conv2d(input_channels, input_channels*2, kernel_size,padding=1,stride=2,padding_mode='reflect')
        self.activation = nn.ReLU() if activation == 'relu' else nn.LeakyReLU(0.2)
        if use_bn:
            self.normalization = nn.InstanceNorm2d(input_channels)
        self.use_bn = use_bn

    def forward(self,x):
        x = self.conv1(x)
        if self.use_bn:
            self.normalization(x)
        x = self.activation(x)
        return x
    

class ExpandingBlock(nn.Module):
    def __init__(self,input_channels,use_bn=True) -> None:
        super(ExpandingBlock, self).__init__()
        self.conv1 = nn.ConvTranspose2d(input_channels, input_channels // 2, kernel_size=3,stride=2,padding=1,output_padding=1)
        if use_bn:
            self.normalization = nn.InstanceNorm2d(input_channels // 2)
        self.use_bn = use_bn
        self.activation = nn.ReLU()

    def forward(self, x):
        x = self.conv1(x)
        if self.use_bn:
            x = self.normalization(x)
        x = self.activation(x)
        return x
    


class FeatureMapBlock(nn.Module):
    def __init__(self, input_channels, output_channels) -> None:
        super(FeatureMapBlock,self).__init__()
        self.conv = nn.Conv2d(input_channels, output_channels,kernel_size=7,padding=3,padding_mode='reflect')

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

class Generator(nn.Module):
    def __init__(self, input_channels,output_channels, hidden_dim=64) -> None:
        super(Generator,self).__init__()
        self.upfeature = FeatureMapBlock(input_channels,hidden_dim)
        self.contract1 = ContractingBlock(hidden_dim)
        self.contract2 = ContractingBlock(hidden_dim * 2)
        res_mult = 4
        self.res0 = ResidualBlock(hidden_dim * res_mult)
        self.res1 = ResidualBlock(hidden_dim * res_mult)
        self.res2 = ResidualBlock(hidden_dim * res_mult)
        self.res3 = ResidualBlock(hidden_dim * res_mult)
        self.res4 = ResidualBlock(hidden_dim * res_mult)
        self.res5 = ResidualBlock(hidden_dim * res_mult)
        self.res6 = ResidualBlock(hidden_dim * res_mult)
        self.res7 = ResidualBlock(hidden_dim * res_mult)
        self.res8 = ResidualBlock(hidden_dim * res_mult)
        self.expand1 = ExpandingBlock(hidden_dim * res_mult)
        self.expand2 = ExpandingBlock(hidden_dim * 2)
        self.downfeature  = FeatureMapBlock(hidden_dim,output_channels)
        self.tanh = nn.Tanh()

    def forward(self, x):
        x0 = self.upfeature(x)
        x1 = self.contract1(x0)
        x2 = self.contract2(x1)
        x3 = self.res0(x2)
        x4 = self.res1(x3)
        x5 = self.res2(x4)
        x6 = self.res3(x5)
        x7 = self.res4(x6)
        x8 = self.res5(x7)
        x9 = self.res6(x8)
        x10 = self.res7(x9)
        x11 = self.res8(x10)
        x12 = self.expand1(x11)
        x13 = self.expand2(x12)
        xn = self.downfeature(x13)
        return self.tanh(xn)
    

class Discriminator(nn.Module):
    def __init__(self, input_channels, hidden_channels=64) -> None:
        super(Discriminator,self).__init__()
        self.upfeature = FeatureMapBlock(input_channels,hidden_channels)
        self.contract1 = ContractingBlock(hidden_channels, False,kernel_size=4,activation='lrelu')
        self.contract2 = ContractingBlock(hidden_channels * 2,kernel_size=4,activation='lrelu')
        self.contract3 = ContractingBlock(hidden_channels * 4,kernel_size=4,activation='lrelu')
        self.conv = nn.Conv2d(hidden_channels*8,1,kernel_size=1)

    def forward(self,x):
        x0 = self.upfeature(x)
        x1 = self.contract1(x0)
        x2 = self.contract2(x1)
        x3 = self.contract3(x2)
        x4 = self.conv(x3)
        return x4