from __future__ import print_function
import argparse
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
import torch.nn.functional as F
from torch.autograd import Variable
import os
import numpy as np

import models.began as began
from folder import ImageFolder

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', required=True, help='cifar10 | lsun | imagenet | folder | lfw ')
parser.add_argument('--dataroot', required=True, help='path to dataset')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=2)
parser.add_argument('--batchSize', type=int, default=16, help='input batch size')
parser.add_argument('--imageSize', type=int, default=32, help='the height / width of the input image to network')
parser.add_argument('--nz', type=int, default=64, help='size of the latent z vector')
parser.add_argument('--size', type=int, default=3, help='size of images generated by G,3:32,4:64,5:128')
parser.add_argument('--gamma', type=float, default=0.5, help='gamma, diversity ratio')
parser.add_argument('--niter', type=int, default=25, help='number of epochs to train for')
parser.add_argument('--lrD', type=float, default=0.00005, help='learning rate for Critic, default=0.00005')
parser.add_argument('--lrG', type=float, default=0.00005, help='learning rate for Generator, default=0.00005')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--cuda'  , action='store_true', help='enables cuda')
parser.add_argument('--ngpu'  , type=int, default=1, help='number of GPUs to use')
parser.add_argument('--netG', default='', help="path to netG (to continue training)")
parser.add_argument('--netD', default='', help="path to netD (to continue training)")
parser.add_argument('--Diters', type=int, default=5, help='number of D iters per each G iter')
parser.add_argument('--experiment', default=None, help='Where to store samples and models')
parser.add_argument('--adam', action='store_true', help='Whether to use adam (default is rmsprop)')
opt = parser.parse_args()
print(opt)

if opt.experiment is None:
    opt.experiment = 'samples'
os.system('mkdir {0}'.format(opt.experiment))

opt.manualSeed = random.randint(1, 10000) # fix seed
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)

cudnn.benchmark = False
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

############# DataLoader ######################
if opt.dataset == 'lsun':
    dataset = dset.LSUN(db_path=opt.dataroot, classes=['bedroom_train'],
                        transform=transforms.Compose([
                            transforms.Scale(opt.imageSize),
                            transforms.CenterCrop(opt.imageSize),
                            transforms.ToTensor(),
                            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                        ]))
elif opt.dataset == 'CelebA':
    dataset = ImageFolder(root=opt.dataroot,
                        transform=transforms.Compose([
                            transforms.Resize(opt.imageSize),
                            transforms.CenterCrop(opt.imageSize),
                            transforms.ToTensor(),
                            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                        ]))

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

ngpu = int(opt.ngpu)
nz = int(opt.nz)

#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)

def L_loss(x, D_x, eta=1):
    if eta==1:
        return torch.mean(torch.abs(x - D_x))
    elif eta==2:
        return torch.mean((x - D_x).pow(2))

netG = began.BEGAN_G(opt.size)
netG.apply(weights_init)
if opt.netG != '': # load checkpoint if needed
    netG.load_state_dict(torch.load(opt.netG))
print(netG)

netD = began.BEGAN_D(opt.size)
netD.apply(weights_init)
if opt.netD != '':
    netD.load_state_dict(torch.load(opt.netD))
print(netD)

input = torch.FloatTensor(opt.batchSize, 3, opt.imageSize, opt.imageSize)
noise = torch.FloatTensor(opt.batchSize, nz)
fixed_noise = torch.FloatTensor(opt.batchSize, nz).normal_(0, 1)

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

# setup optimizer
optimizerD = optim.Adam(netD.parameters(), lr=opt.lrD, betas=(opt.beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=opt.lrG, betas=(opt.beta1, 0.999))

gen_iterations = 0
k = Variable(torch.from_numpy(np.array([0.0])).float(), requires_grad=False)
gamma = Variable(torch.from_numpy(np.array([opt.gamma])).float(), requires_grad=False)
lambda_k = Variable(torch.from_numpy(np.array([0.001])).float(),requires_grad=False)
if opt.cuda:
    k = k.cuda()
    gamma = gamma.cuda()
    lambda_k = lambda_k.cuda()

for epoch in range(opt.niter):
    data_iter = iter(dataloader)
    i = 0
    while i < len(dataloader):
        ############################
        # (1) Update D network
        ###########################
        for p in netD.parameters(): # reset requires_grad
            p.requires_grad = True # they are set to False below in netG update

        # train the discriminator Diters times
        data = data_iter.next()
        i += 1

        # train with real
        real_cpu, _ = data
        netD.zero_grad()
        batch_size = real_cpu.size(0)

        if opt.cuda:
            real_cpu = real_cpu.cuda()
        input.resize_as_(real_cpu).copy_(real_cpu)
        inputv = Variable(input)
        errD_real = netD(inputv)
        L_errD_real = L_loss(inputv, errD_real)
        # train with fake
        noise.resize_(opt.batchSize, nz).normal_(0, 1)
        noisev = Variable(noise, volatile = True) # totally freeze netG
        fake = Variable(netG(noisev).data)
        inputv = fake
        errD_fake = netD(inputv)
        L_errD_fake = L_loss(inputv, errD_fake)
        errD = L_errD_real - k.data[0] * L_errD_fake
        errD.backward()
        optimizerD.step()

        ############################
        # (2) Update G network
        ###########################
        for p in netD.parameters():
            p.requires_grad = False # to avoid computation
        netG.zero_grad()
        # in case our last batch was the tail batch of the dataloader,
        # make sure we feed a full batch of noise
        noise.resize_(opt.batchSize, nz).normal_(0, 1)
        noisev = Variable(noise)
        fake = netG(noisev)
        errG_fake = netD(fake)
        errG = L_loss(fake, errG_fake)
        errG.backward()
        optimizerG.step()
        gen_iterations += 1

        k = k + lambda_k * (gamma * L_errD_real - errG)
        #k = k + lambda_k * (gamma.data[0] * L_errD_real.data[0] - errG.data[0])
        M_global = L_errD_real + torch.abs(lambda_k * L_errD_real - errG)
        print('[%d/%d][%d/%d][%d] L_errD_real: %f L_errD_fake: %f errD: %f errG: %f k:%f M_global: %f'
            % (epoch, opt.niter, i, len(dataloader), gen_iterations,
            L_errD_real.data[0], L_errD_fake.data[0], errD.data[0], errG.data[0],
            k.data[0], M_global.data[0]))

        if gen_iterations % 500 == 0:
            real_cpu = real_cpu.mul(0.5).add(0.5)
            vutils.save_image(real_cpu, '{0}/real_samples.png'.format(opt.experiment))
            fake = netG(Variable(fixed_noise, volatile=True))
            fake.data = fake.data.mul(0.5).add(0.5)
            vutils.save_image(fake.data, '{0}/fake_samples_{1}.png'.format(opt.experiment, gen_iterations))

    # do checkpointing
    torch.save(netG.state_dict(), '{0}/netG_epoch_{1}.pth'.format(opt.experiment, epoch))
    torch.save(netD.state_dict(), '{0}/netD_epoch_{1}.pth'.format(opt.experiment, epoch))
