"""
Train code for shapenet car dataset. Loads the real images from blender , renders them as real image, applies
generated texture to the same mesh and renders them as fake images: trains our whole pipeline on shapenet car dataset
"""

import argparse
import matplotlib.pyplot as plt
import torch
from pytorch3d.io import load_obj
from Utils.utils import weights_init
from torch import nn
from torch import optim
from pytorch3d.renderer import look_at_view_transform, FoVPerspectiveCameras
from pytorch3d.renderer import TexturesUV
from pytorch3d.structures import Meshes
import time
import os
from Utils.utils import get_lights, get_renderer
from Utils.plot import plot_rendered_images, plot_losses
from PIL import Image
from torchvision import transforms
import warnings

warnings.filterwarnings("ignore")

# Analysze time it takes to run the code and memory used
# with clean up inside loop and without it
start_time = time.time()

clean = True

# save path for trained generator
save_path = "Models/generator.pth"
# batch size used for training
batch_size = 1
# spatial size of image generated from generator/ input to Discriminator
image_size = 512

if image_size == 64:
    from Generator.generator64 import Generator
    from Discriminator.discriminator64 import Discriminator
elif image_size == 128:
    from Generator.generator128 import Generator
    from Discriminator.discriminator128 import Discriminator
elif image_size == 512:
    from Generator.generator512 import Generator
    from Discriminator.discriminator512 import Discriminator
else:
    from Generator.generator1024 import Generator
    from Discriminator.discriminator1024 import Discriminator

# num epochs to train
num_epochs = 100
# Learning rate for optimizers
lr = 0.0002
# Beta1 hyperparam for Adam optimizers
beta1 = 0.5
# device to train: gpu/cpu
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# data_path
data_path = "/home/user/data/shapenet_car/02958343"
# num_views: total number of rendered views from blender python api
num_views = 100
folders = os.listdir(data_path)

# parse the arguments for config of Generator and config of Discriminator
parser = argparse.ArgumentParser()
parser.add_argument("--configG", help="configuration file for Generator")
parser.add_argument("--configD", help="configuration file for Discriminator")
args = parser.parse_args()

# netG and netD
print("Initializing Generator...")
netG = Generator(args.configG).to(device)
print(netG)

print("Initializing Discriminator...")
netD = Discriminator(args.configD).to(device)
print(netD)

print("Initializing weights for Generator and Discriminator...")
netG.apply(weights_init)
netD.apply(weights_init)

criterion = nn.BCELoss()
# criterion = nn.BCEWithLogitsLoss()

real_label = 1.
fake_label = 0.

optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))

loss_G = []
loss_D = []

print("Training...")
gpu_usage = []
total_obj = len(folders)
step_size = 360.0 / num_views

for epoch in range(num_epochs):
    # for each epoch iterate over the obj files and train
    print(f"Epoch: {epoch + 1} / {num_epochs}")
    obj_num = 1
    for idx in range(total_obj):
        # load the object
        obj_filename = data_path + "/" + folders[idx] + "/models/" + "bpy_model.obj"
        verts, faces, aux = load_obj(obj_filename, device=device)
        mesh = Meshes([verts.to(device)], [faces.verts_idx.to(device)])
        # load the mesh with uv information and apply generated texture
        faces_uvs = [faces.textures_idx]
        # verts_uvs = mesh.textures._verts_uvs_list
        verts_uvs = [aux.verts_uvs]

        for num_view in range(num_views):
            # load the real image
            image_name = 'models_r_{0:03d}.png'.format(int(num_view * step_size))
            image_path = data_path + "/" + folders[idx] + "/models/rendered_images/" + image_name
            real_image = Image.open(image_path).convert("RGB")
            transform = transforms.Compose([
                transforms.ToTensor(),
                # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
            ])
            # real_batch = transforms.ToTensor()(real_image).unsqueeze_(dim=0).to(device=device)
            real_batch = transform(real_image).unsqueeze_(dim=0).to(device=device)

            # get camera for the same angle
            R, T = look_at_view_transform(2, 30.0, 90 + num_view * step_size)
            cameras = FoVPerspectiveCameras(R=R, T=T, device=device)

            # get lights
            lights = get_lights(device=device)
            # get renderer
            renderer = get_renderer(image_size=image_size, cameras=cameras, lights=lights, device=device)

            # zero grad the netD
            netD.zero_grad()

            label = torch.full((batch_size,), real_label, dtype=torch.float, device=device)
            # forward pass the real batch through D
            output = netD(real_batch).view(-1)
            errD_real = criterion(output, label)
            # calculate gradients for real batch
            errD_real.backward()

            # create one fake texture
            noise = netG.generate_noise(batch_size=1).to(device)
            texture = netG(noise).permute(0, 2, 3, 1)

            # apply fake texture to given mesh with previous configuration
            mesh.textures = TexturesUV(texture.to(device), faces_uvs=faces_uvs, verts_uvs=verts_uvs)

            # will produce fake images of shape [batch_sizeximage_sizeximage_sizex4(RGBA)]
            fake_images = renderer(mesh)
            fake_batch = fake_images[..., :3].to(device).permute(0, 3, 1, 2)
            label.fill_(fake_label)
            # pass fake batch to netD
            output = netD(fake_batch.detach()).view(-1)
            errD_fake = criterion(output, label)
            # calculated the gradient of this fake batch: gradients will accumulate with previous real_batch
            errD_fake.backward()
            # update the discriminator
            optimizerD.step()

            # list some values for plot
            errD = errD_real.item() + errD_fake.item()
            loss_D.append(errD)

            # Train generator
            netG.zero_grad()
            label.fill_(real_label)
            # D got updated, make another pass to create computation graph and then to calculate the gradients
            output = netD(fake_batch).view(-1)
            # calculate loss
            errG = criterion(output, label)
            # calculate gradients
            errG.backward()
            # update generator
            optimizerG.step()

            # some book keeping
            loss_G.append(errG.item())

            # Output training stats
            print('[%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f' % (obj_num, total_obj, errD, errG.item()))

        obj_num += 1

        # clean up
        if clean:
            del mesh, cameras, renderer, verts_uvs, faces_uvs, fake_images, texture

# save the trained generator
torch.save(netG.state_dict(), save_path)

end_time = time.time()
elapsed_minutes = (end_time - start_time) / 60
print(f"Time: {round(elapsed_minutes, 2)} Minutes")

# plot losses
plot_losses(loss_G, loss_D)

# generate one texture image
with torch.no_grad():
    noise = netG.generate_noise(batch_size=1).to(device)
    texture = netG(noise).permute(0, 2, 3, 1)

plt.imshow(texture.squeeze().cpu().numpy())
plt.title("sample texture")
plt.show()

obj_filename = data_path + "/" + folders[0] + "/models/" + "bpy_model.obj"
# mesh = load_objs_as_meshes([obj_filename], device=device)
verts, faces, aux = load_obj(obj_filename, device=device)
mesh = Meshes([verts.to(device)], [faces.verts_idx.to(device)])
# load the mesh with uv information and apply generated texture
faces_uvs = [faces.textures_idx]
# verts_uvs = mesh.textures._verts_uvs_list
verts_uvs = [aux.verts_uvs]

R, T = look_at_view_transform(dist=2, elev=30, azim=90)
cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
# get lights
lights = get_lights(device=device)
# get renderer
renderer = get_renderer(image_size=image_size, cameras=cameras, lights=lights, device=device)
mesh.textures = TexturesUV(texture.to(device), faces_uvs=faces_uvs, verts_uvs=verts_uvs)

# render images: will produce a tensor of shape: [batch_size, image_size, image_size, 4(RGBA)]
image = renderer(mesh, cameras=cameras, lights=lights)

plot_rendered_images(image)
