import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
import torch.nn.functional as F
# Initialize the MSELoss function
criterion1 = nn.MSELoss()
# Initialize the L1Loss function
criterion2 = nn.L1Loss()    
real_label = Variable(torch.cuda.FloatTensor(1).fill_(1.0), requires_grad = False)
fake_label = Variable(torch.cuda.FloatTensor(1).fill_(0.0), requires_grad = False)
def cal_loss_Gan(D, real, fake):
    '''
    input:
        D--Discriminator
        real--X from X domain  or Y from Y domain 
        fake--F(Y) generated by using Y from Y domain or G(X) generated by using X from X domain
    '''
    pred_real = D(real)
    pred_fake = D(fake.detach())
    loss_D_real = criterion1(pred_real, real_label)
    loss_D_fake = criterion1(pred_fake, fake_label)
    loss_D = 0.5 * (loss_D_fake + loss_D_real) 
    return loss_D

def cal_loss_Cycle(net, real, fake):
    '''
    input:
        net:
            G--Generator which generate image from X domain to Y domain
            or F--Generator which generate image from Y domain to X domain
        real--X from X domain  or Y from Y domain 
        fake--F(Y) generated by using Y from Y domain or G(X) generated by using X from X domain
    return: Cycle loss
    '''
    loss_Cycle = criterion2(real, net(fake))
    return loss_Cycle

def weights_init(m):
    
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)

class Discriminator(nn.Module):
    
    def __init__(self,input_nc):
        super(Discriminator, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1,padding=1)
        self.conv2 = nn.Conv2d(16, 64, kernel_size=3, stride=1,padding=1)
        self.conv3 = nn.Conv2d(64, 256, kernel_size=3, stride=1,padding=1)
        self.conv4 = nn.Conv2d(256, 256, kernel_size=3, stride=1,padding=1)
        self.relu = nn.ReLU(inplace=True)
        
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(256 * 56 * 56, 100)
        self.fc2 = nn.Linear(100, 1)
        self.features = nn.Sequential(self.conv1, self.relu, self.conv2,self.relu, self.maxpool,
                                      self.conv3,self.relu, self.conv4,self.relu, self.maxpool )
        self.determine = nn.Sequential(self.fc1,self.relu,self.fc2)
        
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.determine(x)

        return x

class Generator(nn.Module):

    def __init__(self, D, C=64):
        super(Generator, self).__init__()
        self.D = D
        self.conv = nn.ModuleList()
        self.conv.append(nn.Conv2d(3, C, 3, padding=(1,1)))
        nn.init.kaiming_normal_(self.conv[0].weight.data)

        for k in range(D):
            self.conv.append(nn.Conv2d(C,C,3,padding=(1,1)))
            nn.init.kaiming_normal_(self.conv[-1].weight.data)
            
        self.conv.append(nn.Conv2d(C,3,3,padding=(1,1)))
        
        self.bn = nn.ModuleList()
        for k in range(D):
            self.bn.append(nn.BatchNorm2d(C, C))
            nn.init.constant_(self.bn[k].weight.data, 1.25 * np.sqrt(C))
        self.relu = nn.ReLU(inplace = True)
        
    def forward(self, x):
        D = self.D
        h = self.relu(self.conv[0](x))
        for k in range(D):
            h = self.relu(self.bn[k](self.conv[k+1](h)))
            
        y = self.conv[D+1](h) + x

        return y