import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import utils
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
from torch.autograd import Variable
import numpy as np
import os
import argparse
import network
import time
from datasets import CUBDataset



parser = argparse.ArgumentParser(description='PyTorch GAN for MNIST')
parser.add_argument('--epochs', type=int, default=1000, metavar='N',
                    help='number of epochs to train (default: 1000)')
parser.add_argument('--batch-size', type=int, default=100, metavar='N',
                    help='input batch size for training (default: 100)')
#parser.add_argument('--GPU', type=int, default='CPU', metavar='N',
#                    help='mode (default: CPU)')
args = parser.parse_args()



batch_size = args.batch_size
data_transform = transforms.Compose([
        transforms.Scale((64,64)),
	transforms.ToTensor(),
        transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
    ])
trainset = CUBDataset('/beegfs/jn1664/CUB_200_2011', train=True, transform=data_transform)
testset = CUBDataset('/beegfs/jn1664/CUB_200_2011', train=False, transform=data_transform)
#trainset = dsets.MNIST(root='.', train=True, download=True, transform=data_transform)
#testset = dsets.MNIST(root='.', train=False, download=True, transform=data_transform)

train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=0)
test_loader  = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=0)


def showImg(img, epoch):
    num_img = img.shape[0]
    img = img.view(num_img,3, 64,64)
    img = utils.make_grid(img)
    # unnormalize : [-0.424213, 2.82149] with mean 0 => [0,1] with mean 0.1308
    #img = img / 3.2457 + 0.1307
    img = img / 2 + 0.5
    npimg = img.numpy()
    plt.imsave('img_gan/{}.png'.format(epoch), np.transpose(npimg, (1, 2, 0)))
    #grid_img = plt.imshow(np.transpose(npimg, (1, 2, 0)))
    #plt.savefig('img_gan/{}.png'.format(epoch))
    #plt.show(grid_img)




g = network.Generator(100, 4096*3)
d = network.Discriminator(4096*3, 1)
g.cuda()
d.cuda()

criterion = nn.BCELoss()
lr_g = 0.0002
lr_d = 0.01
g_optimizer = torch.optim.Adam(g.parameters(), lr=lr_g)
d_optimizer = torch.optim.SGD(d.parameters(), lr=lr_d)



def train(epoch):
    if os.path.exists('model_0/model_1100_G.pth'):
	g.load_state_dict(torch.load('model_0/model_1100_G.pth')) 
    if os.path.exists('model_0/model_1100_D.pth'):
        d.load_state_dict(torch.load('model_0/model_1100_D.pth'))


    g.train()
    d.train()


    for batch_idx, (data,target) in enumerate(train_loader,1):
	num_target = data.shape[0]


	#prepare-1. real data
        data = Variable(data.cuda())
	#prepare-2. description
	text = Variable(target.cuda())
	#prepare-3. noise data
	noise = Variable(torch.randn(num_target, 100).cuda())
	#prepare-4. real label
	real_label = Variable(torch.ones(num_target).cuda())
	#prepare-5. fake label
	fake_label = Variable(torch.zeros(num_target).cuda())


        #D-1. discriminate for real images
        d_real_output = d(data, text)
        #D-2. compute real loss
        d_real_loss = criterion(d_real_output, real_label)
        #D-3. discriminate for fake number images generated by G
        d_fake_output = d(g(noise, text, num_target), text)
        #D-4. compute total loss and optimize
	d.zero_grad()
        d_fake_loss = criterion(d_fake_output, fake_label)
	d_loss = d_real_loss + d_fake_loss
        d_loss.backward()
        d_optimizer.step()



	#G-1 compute loss and optimize
        fake_img = g(noise, text, num_target)#generate a fake image
        d_output = d(fake_img, text)#get results from d
	g.zero_grad()
        g_loss = criterion(d_output, real_label)#compute loss
        g_loss.backward()#optimize
        g_optimizer.step()




        if batch_idx % 10 == 0:
            print('Generator Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), g_loss.data[0]))
            print('Descriminator Train Epoch: {} [{}/{} ({:.0f}%)]\tReal Loss: {:.6f}\tFake Loss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), d_real_loss.data[0], d_fake_loss.data[0]))


	if epoch % 10 == 0:
            # show images
            showImg(fake_img.data.cpu(), epoch)




num_epochs = args.epochs
for epoch in range(num_epochs):
    start_t = time.time()
    train(epoch)
    end_t = time.time()
    print('''[Total Train time]: %.2fsec'''% (end_t - start_t))


    if epoch % 100 == 0:
	model_file = 'model/model_' + str(epoch) + '_D.pth'
	torch.save(d.state_dict(), model_file)
	model_file = 'model/model_' + str(epoch) + '_G.pth'
        torch.save(g.state_dict(), model_file)

