import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from NetworkArchitecture import Discriminator, Generator, initialize_weights
import time
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

#hyper params
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
LEARNING_RATE = 2e-4
BATCH_SIZE = 32
IMAGE_SIZE = 64
CHANNELS_IMG = 1

NOISE_DIM = 10
Z_DIM = 100
NUM_EPOCHS = 5
FEATURES_DISC = 64
FEATURES_GEN = 64
LOAD_PARAMS = False


#Sample checkpoint
'''
checkpoint = {"gen_state_dict": gen.state_dict(), 'gen_opt': opt_gen.state_dict()
        "disc_state_dict": disc.state_dict(), 'disc_opt':opt_disc.state_dict()}
'''
def save_checkpoint(state, filename="my_checkpoint.pth.tar"):
    print("=> Saving Checkpoint")
    torch.save(state, filename)

def load_checkpoint(filename="my_checkpoint.pth.tar"):
    print("=> Loading Checkpoint")
    checkpoint = torch.load(filename)
    gen.load_state_dict(checkpoint['gen_state_dict'])
    opt_gen.load_state_dict(checkpoint['gen_opt'])
    disc.load_state_dict(checkpoint['disc_state_dict'])
    disc_opt.load_state_dict(checkpoint['disc_opt'])

transforms = transforms.Compose(
    [
        transforms.Resize(IMAGE_SIZE),
        transforms.ToTensor(),
        transforms.Normalize(
            [0.5 for _ in range(CHANNELS_IMG)], [0.5 for _ in range(CHANNELS_IMG)])
    ]
)

dataset = datasets.MNIST(root="dataset/", train=True, transform=transforms, download=True)
#dataset = datasets.ImageFolder(root="audio_images_cropped", transform=transforms, train=False)
loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
gen = Generator(Z_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)
disc = Discriminator(CHANNELS_IMG, FEATURES_DISC).to(device)

initialize_weights(gen)
initialize_weights(disc)

opt_gen = optim.Adam(gen.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))
opt_disc = optim.Adam(disc.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))
criterion = nn.BCELoss()

fixed_noise = torch.randn(32, Z_DIM, 1, 1).to(device)
writer_real = SummaryWriter(f"musical_logs/real")
writer_fake = SummaryWriter(f"musical_logs/fake")
step = 0

gen.train()
disc.train()

if(LOAD_PARAMS):
    load_checkpoint("my_checkpoint.pth.tar")

for epoch in range(NUM_EPOCHS):

    
    if (epoch % 2 == 0):
        checkpoint = {"gen_state_dict": gen.state_dict(), 'gen_opt': opt_gen.state_dict(),
        "disc_state_dict": disc.state_dict(), 'disc_opt':opt_disc.state_dict()}
        save_checkpoint(checkpoint)


    for batch_idx, (real, _) in enumerate(loader):
        real = real.to(device)
        noise = torch.randn((BATCH_SIZE, Z_DIM, 1, 1)).to(device)
        fake = gen(noise)

        #Train Discriminator max log(D(x)) + log(1 - D(G(z)))
        disc_real = disc(real).reshape(-1) 
        loss_disc_real = criterion(disc_real, torch.ones_like(disc_real))
        disc_fake = disc(fake).reshape(-1)
        loss_disc_fake = criterion(disc_fake, torch.zeros_like(disc_fake))
        loss_disc = (loss_disc_real + loss_disc_fake) / 2
        disc.zero_grad()
        loss_disc.backward(retain_graph=True)
        opt_disc.step()

        #Train Generator max log(D(G(z)))
        output = disc(fake).reshape(-1)
        loss_gen = criterion(output, torch.ones_like(output))
        gen.zero_grad()
        loss_gen.backward()
        opt_gen.step()

        

        #Print losses occasionally and print to tensorboard
        if batch_idx % 100 == 0:
            print(
                f"Epoch [{epoch}/{NUM_EPOCHS}] Batch {batch_idx}/{len(loader)} \
                Loss D: {loss_disc:.4f}, loss G: {loss_gen: .4f}"
            )

            with torch.no_grad():
                fake = gen(fixed_noise)

                if batch_idx % 500 == 0:
                    images = fake.cpu().numpy()
                    final_image = np.squeeze(images[0]) * 255
                    im = Image.fromarray(final_image)


                    if(im.mode != "L"):
                        im = im.convert('L')
                    im.save("test.png")
                    plt.imshow(final_image)
                    plt.show()

                img_grid_real = torchvision.utils.make_grid(
                    real[:32], normalize=True
                )
                img_grid_fake = torchvision.utils.make_grid(
                    fake[:32], normalize=True
                )

                writer_real.add_image("Real", img_grid_real, global_step=step)
                writer_fake.add_image("Fake", img_grid_fake, global_step=step)

            step += 1


