import torch.nn as nn
from utils import *



class Encoder(nn.Module):
    def  __init__(self,input_channels, use_bn=True, use_dropout = False):
        super(Encoder, self).__init__()
        self.conv1 = nn.Conv2d(input_channels, 2 * input_channels, kernel_size=3)
        self.conv2 = nn.Conv2d(input_channels * 2, 2 * input_channels, kernel_size = 3)
        self.activation = nn.ReLU()
        self.pooling = nn.MaxPool2d(kernel_size=2, stride=2)
        if use_bn:
            self.bn = nn.BatchNorm2d(input_channels * 2)
        self.use_bn = use_bn
        if use_dropout:
            self.dropout = nn.Dropout()
        self.use_dp = use_dropout

    def forward(self, x):
        x = self.conv1(x)
        if self.use_bn:
            x = self.bn(x)
        if self.use_dp:
            x = self.dropout(x)
        x = self.activation(x)
        x = self.conv2(x)
        if self.use_bn:
            x = self.bn(x)
        if self.use_dp:
            x = self.dropout(x)
        x = self.activation(x)
        x = self.pooling(x)
        return x
    

class Decoder(nn.Module):
    def __init__(self, input_channels, use_bn=True, use_dropout = False):
        super(Decoder, self).__init__()
        self.upsample = nn.Upsample(scale_factor = 2, mode='bilinear', align_corners=True)
        self.conv1 = nn.Conv2d(input_channels, input_channels // 2, kernel_size = 2)
        self.conv2 = nn.Conv2d(input_channels , input_channels //2 , kernel_size = 3,padding=1)
        self.conv3 = nn.Conv2d(input_channels //2, input_channels // 2, kernel_size = 2,padding=1)
        self.activation = nn.ReLU()
        if use_bn:
            self.bn = nn.BatchNorm2d(input_channels // 2)
        self.use_bn = use_bn
        if use_dropout:
            self.dropout = nn.Dropout()
        self.use_dp = use_dropout
        

    def forward(self, x, skip_con_x):
        x = self.upsample(x)
        x = self.conv1(x)
        skip_con_x = crop(skip_con_x, x.shape)
        x = torch.cat([x, skip_con_x], axis = 1)
        x = self.conv2(x)
        if self.use_bn:
            x = self.bn(x)
        if self.use_dp:
            x = self.dropout(x)
        x = self.activation(x)
        x = self.conv3(x)
        if self.use_bn:
            x = self.bn(x)
        if self.use_dp:
            x = self.dropout(x)
        x = self.activation(x)
        return x
        

class FeatureMapBlock(nn.Module):
    def __init__(self, input_channels, output_channels):
        super(FeatureMapBlock, self).__init__()
        self.conv = nn.Conv2d(input_channels, output_channels, kernel_size=1)
    def forward(self, x):
        x = self.conv(x)
        return x
    

class UNet(nn.Module):
    def __init__(self, input_channels, output_channels, hidden_channels = 64):
        super(UNet, self).__init__()
        self.upfeature = FeatureMapBlock(input_channels, hidden_channels)
        self.encoder1 = Encoder(hidden_channels, use_dropout=True)
        self.encoder2 = Encoder(hidden_channels * 2,use_dropout=True)
        self.encoder3 = Encoder(hidden_channels * 4, use_dropout=True)
        self.encoder4 = Encoder(hidden_channels * 8)
        self.encoder5 = Encoder(hidden_channels * 16)
        self.encoder6 = Encoder(hidden_channels * 32)
        self.decoder1 = Decoder(hidden_channels * 64)
        self.decoder2 = Decoder(hidden_channels * 32)
        self.decoder3 = Decoder(hidden_channels * 16)
        self.decoder4 = Decoder(hidden_channels * 8)
        self.decoder5 = Decoder(hidden_channels * 4 )
        self.decoder6 = Decoder(hidden_channels * 2)
        self.downfeature = FeatureMapBlock(hidden_channels , output_channels)
        self.sigmoid = nn.Sigmoid()


    def forward(self, x):
        x0 = self.upfeature(x)
        x1 = self.encoder1(x0)
        x2 = self.encoder2(x1)
        x3 = self.encoder3(x2)
        x4 = self.encoder4(x3)
        x5 = self.encoder5(x4)
        x6 = self.encoder6(x5)
        x7 = self.decoder1(x6, x5)
        x8 = self.decoder2(x7, x4)
        x9 = self.decoder3(x8, x3)
        x10 = self.decoder4(x9, x2)
        x11 = self.decoder5(x10, x1)
        x12 = self.decoder6(x11, x0)
        xn = self.downfeature(x12)

        return self.sigmoid(xn)
    

class Discriminator(nn.Module):
    def __init__(self, input_channels, hidden_channels=8) -> None:
        super(Discriminator, self).__init__()
        self.upfeature = FeatureMapBlock(input_channels,hidden_channels)
        self.encoder1 = Encoder(hidden_channels, use_bn=False)
        self.encoder2 = Encoder(hidden_channels*2)
        self.encoder3 = Encoder(hidden_channels*4)
        self.encoder4 = Encoder(hidden_channels*8)
        self.final = nn.Conv2d(hidden_channels * 16,1,kernel_size=1)

    def forward(self, x,y):
        x = torch.cat([x,y],axis=1)
        x0 = self.upfeature(x)
        x1 = self.encoder1(x0)
        x2 = self.encoder2(x1)
        x3 = self.encoder3(x2)
        x4 = self.encoder4(x3)
        xn = self.final(x4)

        return xn