import argparse

from autoencoder import Autoencoder

from copy import deepcopy

from generator import Generator

import os

from torchvision import transforms, datasets

import torchvision

from torchvision.utils import save_image

from torch.distributions import normal, uniform

from torch.utils.data import DataLoader

from transporter import Transporter

from utils import *

torch.set_default_tensor_type('torch.cuda.FloatTensor')

parser = argparse.ArgumentParser()

parser.add_argument('--dataset', default='mnist',

help='dataset name (default: mnist)')

parser.add_argument('--folder', default='.',

help='folder name to put everything in')

parser.add_argument('--steps_a', type=int, default=15000,

help='number of steps to take (default: 15000)')

parser.add_argument('--batchsize_a', type=int, default=128,

help='input batch size for autoencoder (default: 128)')

parser.add_argument('--conv', type=bool, default=False,

help='Is your autoencoder convolutional (default: false)')

parser.add_argument('--dim', type=int, default=30,

help='dimension of autoencoder (default: 30)')

parser.add_argument('--load_a', type=bool, default=False,

help='Whether to load the autoencoder (default: False)')

parser.add_argument('--batchsize_l', type=int, default=128,

help='input batch size for the latent space model (default: 128)')

parser.add_argument('--latent_distr', default='uniform',

help='input batch size for the latent space model (default: uniform)')

parser.add_argument('--load_l', type=bool, default=False,

help='Whether to load the latent space model (default: False)')

parser.add_argument('--steps_l', type=int, default=15000,

help='number of steps your latent space model takes (default: 15000)')

parser.add_argument('--model', default='generator',

help='Is your latent space model a transporter or a generator (default: generator)')

args = parser.parse_args()

import pudb; pudb.set_trace()

DATASET = args.dataset

FOLDER = args.folder

AE_STEPS = args.steps_a

BATCH_SIZE = args.batchsize_a

CONV = args.conv

DIM = args.dim

AE_LOAD = args.load_a

BATCH_SIZE_GEN = args.batchsize_l

DISTR = args.latent_distr

GEN_LOAD = args.load_l

STEPS = args.steps_l

MODEL = args.model

# Load the right dataset

if DATASET == 'mnist':

mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())

mnist_trainloader = DataLoader(mnist_trainset, batch_size=BATCH_SIZE, shuffle=True)

mnist_testset = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())

mnist_testloader = DataLoader(mnist_testset, batch_size=BATCH_SIZE, shuffle=True)

x_train = unload(mnist_trainloader)

x_test = unload(mnist_testloader)

shape = (1, 28, 28)

elif DATASET == 'fashion_mnist':

fashion_trainset = datasets.FashionMNIST(root='./fashion_data', train=True, download=True, transform=transforms.ToTensor())

fashion_trainloader = DataLoader(fashion_trainset, batch_size=BATCH_SIZE, shuffle=True)

fashion_testset = datasets.FashionMNIST(root='./fashion_data', train=False, download=True, transform=transforms.ToTensor())

fashion_testloader = DataLoader(fashion_testset, batch_size=BATCH_SIZE, shuffle=True)

x_train = unload(fashion_trainloader)

x_test = unload(fashion_testloader)

shape = (1, 28, 28)

elif DATASET == 'cifar10':

cifar10_trainset = datasets.CIFAR10(root='./cifar_data', train=True, download=True, transform=transforms.ToTensor())

cifar10_trainloader = DataLoader(cifar10_trainset, batch_size=BATCH_SIZE, shuffle=True)

cifar10_testset = datasets.CIFAR10(root='./cifar_data', train=False, download=True, transform=transforms.ToTensor())

cifar10_testloader = DataLoader(cifar10_testset, batch_size=BATCH_SIZE, shuffle=True)

x_train = unload(cifar10_trainloader)

x_test = unload(cifar10_testloader)

shape = (3, 32, 32)

elif DATASET == 'faces':

transform = transforms.Compose([

transforms.CenterCrop(140),

transforms.Resize(64),

transforms.ToTensor(),

])

dset = torchvision.datasets.ImageFolder('celebA', transform)

dset_size = len(dset)

#print("NOTE!!!! IVE CHANED THE DATASET SIZE!!!")

train_size = 7*dset_size//10 #99*dset_size//100

train_set, test_set = torch.utils.data.random_split(dset, (train_size, dset_size-train_size))

train_loader = torch.utils.data.DataLoader(train_set, batch_size=128, shuffle=True)

test_loader = torch.utils.data.DataLoader(test_set, batch_size=128, shuffle=True)

shape = (3, 64, 64)

else:

raise NotImplementedError

# Train the autoencoder

ae = Autoencoder(shape, DIM, FOLDER, BATCH_SIZE, CONV)

if AE_LOAD:

ae.load_weights("autoencoder")

else:

if "x_train" in locals():

ae.train(AE_STEPS, x_train, x_test, lr=0.003)

else:

ae.train_iter(AE_STEPS, train_loader, test_loader, lr=0.003)

# Prepare the Latent Space Model

if not 'x_test' in locals():

# print ("EVENTUALLY CHANGE THIS BACK AS WELL!!!")

x_test = np.load("faces.npy")

# x_test = unload(test_loader)

encodings = ae.encode(x_test)

if MODEL == 'transporter':

model = Transporter(encodings, DISTR, FOLDER, BATCH_SIZE_GEN)

elif MODEL == 'generator':

model = Generator(encodings, DISTR, FOLDER, BATCH_SIZE_GEN) # I Could try L2 Loss instead of L1?

else:

raise NotImplementedError

# Train the Latent Space Model

if GEN_LOAD:

model.load_weights(MODEL)

else:

model.train(STEPS, lr=0.001) # I should try adjusting the learning rate?

#model.train(STEPS//2, lr=0.0003)

#model.train(STEPS//2, lr=0.0001)

# Display Results

fake_distr = model.generate(batches=1)

fake_img = ae.decode(fake_distr)

fake_img = np.reshape(fake_img, ((BATCH_SIZE_GEN,) + shape))

save_image(torch.Tensor(fake_img[0:64]), os.path.join(FOLDER, "final.png"))

# Save Images in a file for later evaluation

eval_distr = model.generate(batches=10000//BATCH_SIZE_GEN)

eval_img = ae.decode(eval_distr)

np.save("{}/distribution.npy".format(FOLDER), eval_distr)

np.save("{}/images.npy".format(FOLDER), eval_img)

#channels_last = np.rollaxis(fake_img[0:16], 1, 4)

#display_img(channels_last, columns=4)