#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Tue Jul 25 09:37:16 2017

@author: ra
"""

import argparse
import os
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
from torchvision.datasets import CIFAR10
"""
WGAN 相比于DCGAN 的修改：
1. 判别器最后一层去掉sigmoid                                       # 回归问题,而不是二分类概率
2. 生成器和判别器的loss不取log                                      # Wasserstein 距离
3. 每次更新判别器的参数之后把它们的绝对值截断到不超过一个固定常数c          #Ｗ距离－＞Ｌ连续－＞数值稳定
4. 不要用基于动量的优化算法（包括momentum和Adam），推荐RMSProp，SGD也行　 

GAN 两大问题的解释：
collapse mode ->KL 散度不对称
数值不稳定 -> KL散度和JS散度优化方向不一样
"""


class Config:
    lr=0.0002
    nz=100# 噪声维度
    imageSize=64
    nc=3# 图片三通道
    ngf=64 #生成图片
    ndf=64 #判别图片
    gpuids=None
    beta1=0.5
    batchsize=32
    max_epoch=100# =1 when debug
    workers=4
    ngpu = 0
    cuda=False
    outf='/home/ra/pytorch-CycleGAN-and-pix2pix/datasets/genPig' 
    clamp_num=0.01

    
    
opt=Config()
print(opt)

# 数据加载和预处理 
dataset=CIFAR10(root='/home/ra/pytorch-CycleGAN-and-pix2pix/datasets/',train=False,\
                transform=transforms.Compose(\
                                             [transforms.Scale(opt.imageSize) ,
                                              transforms.ToTensor(),
                                              transforms.Normalize([0.5]*3,[0.5]*3)
                                              
                                             ]))
# 

dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchsize,
                                         shuffle=True, num_workers=int(opt.workers))




# custom weights initialization called on netG and netD
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)


class _netG(nn.Module):
    def __init__(self, ngpu):
        super(_netG, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d(     opt.nz, opt.ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(opt.ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(opt.ngf * 8, opt.ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(opt.ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d(opt.ngf * 4, opt.ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(opt.ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d(opt.ngf * 2,     opt.ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(opt.ngf),
            nn.ReLU(True),
            # state size. (ngf) x 32 x 32
            nn.ConvTranspose2d(    opt.ngf,      opt.nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # state size. (nc) x 64 x 64
        )

    def forward(self, input):
        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
        else:
            output = self.main(input)
        return output


netG = _netG(opt.ngpu)
netG.apply(weights_init)

class _netD(nn.Module):
    def __init__(self, ngpu):
        super(_netD, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is (nc) x 64 x 64
            nn.Conv2d(opt.nc,opt. ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf) x 32 x 32
            nn.Conv2d(opt.ndf, opt.ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(opt.ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*2) x 16 x 16
            nn.Conv2d(opt.ndf * 2, opt.ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(opt.ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*4) x 8 x 8
            nn.Conv2d(opt.ndf * 4, opt.ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(opt.ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*8) x 4 x 4
            nn.Conv2d(opt.ndf * 8, 1, 4, 1, 0, bias=False),
            # remove  nn.Sigmoid()
        )

    def forward(self, input):
        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
        else:
            output = self.main(input)
        output=output.mean(0)

        return output.view(1)


netD = _netD(opt.ngpu)
netD.apply(weights_init)


criterion = nn.BCELoss()

input = torch.FloatTensor(opt.batchsize, 3, opt.imageSize, opt.imageSize)
noise = torch.FloatTensor(opt.batchsize, opt.nz, 1, 1)
fixed_noise = torch.FloatTensor(opt.batchsize, opt.nz, 1, 1).normal_(0, 1)
label = torch.FloatTensor(opt.batchsize)
real_label = 1
fake_label = 0

if opt.cuda:
    netD.cuda()
    netG.cuda()
    criterion.cuda()
    input, label = input.cuda(), label.cuda()
    noise, fixed_noise = noise.cuda(), fixed_noise.cuda()

input = Variable(input)
label = Variable(label)
noise = Variable(noise)
fixed_noise = Variable(fixed_noise)

# modify optimizer
optimizerD = optim.RMSprop(netD.parameters(), lr=opt.lr)
optimizerG = optim.RMSprop(netG.parameters(), lr=opt.lr)

positive=torch.FloatTensor([1])
negtive=-1*positive

for epoch in range(opt.max_epoch):
    for i, data in enumerate(dataloader, 0):
        ############################
       
        ###########################
        # train with real
        netD.zero_grad()
        real_cpu, _ = data
        batch_size = real_cpu.size(0)
        input.data.resize_(real_cpu.size()).copy_(real_cpu)
        label.data.resize_(batch_size).fill_(real_label)

        errD_real = netD(input)
        #errD_real = criterion(output, label)
        errD_real.backward(positive)
        D_x = output.data.mean()

        # train with fake
        noise.data.resize_(batch_size, opt.nz, 1, 1)
        noise.data.normal_(0, 1)
        fake = netG(noise)
        label.data.fill_(fake_label)
        errD_fake = netD(fake.detach())
       # errD_fake = criterion(output, label)
        errD_fake.backward(negtive)
        D_G_z1 = output.data.mean()
        errD = errD_real - errD_fake
        optimizerD.step()
        for parm in netD.parameters():
            parm.data.clamp_(-opt.clamp_num,opt.clamp_num) ### 只有判别器需要 截断参数
        

        ############################
        # (2) Update G network
        ###########################
        if torch.rand(1)[0]>0.8:
            netG.zero_grad()
            label.data.fill_(real_label)  # fake labels are real for generator cost
            output = netD(fake)
            #errG = criterion(output, label)
            output.backward(positive)
            D_G_z2 = output.data.mean()
            optimizerG.step()

        print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'
              % (epoch+1, opt.max_epoch, i, len(dataloader),
                 errD.data[0], errG.data[0], D_x, D_G_z1, D_G_z2))
        if i % 100 == 0:
            vutils.save_image(real_cpu,
                    '%s/real_samples_epoch_%03d.png' % (opt.outf, epoch),
                    normalize=True)
            fake = netG(fixed_noise)
            vutils.save_image(fake.data,
                    '%s/fake_samples_epoch_%03d.png' % (opt.outf, epoch),
                    normalize=True)

    # do checkpointing
#    torch.save(netG.state_dict(), '%s/netG_epoch_%d.pth' % (opt.outf, epoch))
#    torch.save(netD.state_dict(), '%s/netD_epoch_%d.pth' % (opt.outf, epoch))
