from discriminator import CycleGANDiscriminator
from generator import CycleGANGenerator
from imagebuffer import ImageBuffer
from cyclegan_dataset import CycleGANTestDataset
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from PIL import Image

import os
import glob
from torch.utils.tensorboard import SummaryWriter
from cyclegan_dataset import CycleGANDataset
from torch.utils.data.dataloader import DataLoader
import torchvision
import torchvision.utils as vutils

class GAN:
    def __init__(self):
        # Use the gpu
        self.device = torch.device('cuda')

        # Create the generators and discriminators
        self.generator_A = CycleGANGenerator(3, 3, 64).to(self.device)
        self.generator_B = CycleGANGenerator(3, 3, 64).to(self.device)
        self.discriminator_A = CycleGANDiscriminator(3, 64).to(self.device)
        self.discriminator_B = CycleGANDiscriminator(3, 64).to(self.device)

        # Print the networks
        print(self.generator_A)
        print(self.generator_B)
        print(self.discriminator_A)
        print(self.discriminator_B)

        # Initialize the weights of all networks
        self.generator_A.apply(self.init_weights)
        self.generator_B.apply(self.init_weights)
        self.discriminator_A.apply(self.init_weights)
        self.discriminator_B.apply(self.init_weights)

        # Create the optimizers for all the networks
        self.generator_A_optimizer = optim.Adam(self.generator_A.parameters(), lr = 0.0002, betas = (0.5, 0.999))
        self.generator_B_optimizer = optim.Adam(self.generator_B.parameters(), lr = 0.0002, betas = (0.5, 0.999))
        self.discriminator_A_optimizer = optim.Adam(self.discriminator_A.parameters(), lr = 0.0002, betas = (0.5, 0.999))
        self.discriminator_B_optimizer = optim.Adam(self.discriminator_B.parameters(), lr = 0.0002, betas = (0.5, 0.999))

        # Create learning rate schedulers for all the optimizers
        self.generator_A_scheduler = optim.lr_scheduler.LambdaLR(self.generator_A_optimizer, lr_lambda = self.schedule_rate)
        self.generator_B_scheduler = optim.lr_scheduler.LambdaLR(self.generator_B_optimizer, lr_lambda = self.schedule_rate)
        self.discriminator_A_scheduler = optim.lr_scheduler.LambdaLR(self.discriminator_A_optimizer, lr_lambda = self.schedule_rate)
        self.discriminator_B_scheduler = optim.lr_scheduler.LambdaLR(self.discriminator_B_optimizer, lr_lambda = self.schedule_rate)

        # Create the buffers to store history of images generated by the generators
        self.generator_A_buffer = ImageBuffer(buffer_size = 50)
        self.generator_B_buffer = ImageBuffer(buffer_size = 50)

        # Define the loss criterions
        self.cycle_loss = nn.L1Loss()
        self.gan_loss = nn.MSELoss()

        # Get the dataset and dataloaders
        self.dataset = CycleGANDataset(base_dir = '/home/paurvi/CycleGAN/datasets/summer2winter_yosemite')
        self.dataloader = DataLoader(self.dataset, batch_size = 1, num_workers = 2)

        # Writers to tensorboard
        self.steps = 0
        self.writer = SummaryWriter(comment = 'cyclegan_cityscapes- 001 changed CycleGANDiscriminator leaakyRelu slope from 0.2 to 0.5')
    
    def save_network(self, epoch):
        """ Save all the networks
        """
        path = os.path.join(os.getcwd(), 'models', 'cyclegan')

        torch.save(self.generator_A.state_dict(), '%s/generator_A_epoch_%d.pth' % (path, epoch))
        torch.save(self.generator_B.state_dict(), '%s/generator_B_epoch_%d.pth' % (path, epoch))
        torch.save(self.discriminator_A.state_dict(), '%s/discriminator_A_epoch_%d.pth' % (path, epoch))
        torch.save(self.discriminator_B.state_dict(), '%s/discriminator_B_epoch_%d.pth' % (path, epoch))

    def schedule_rate(self, epoch):
        epoch = max(0, epoch - 100)
        return max(0.0, 1.0 - epoch / 100)

    def set_grads(self, networks, grad):
        for network in networks:
            for param in network.parameters():
                param.requires_grad = grad

    def init_weights(self, m):
        """ Initialize weights for the network
        """
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
            nn.init.normal_(m.weight.data, mean = 0.0, std = 0.02)
            if hasattr(m, 'bias') and m.bias is not None:
                nn.init.constant_(m.bias.data, val = 0.0)
    
    def optimize(self, imageA, imageB, epoch, iteration):
        """ imageA : sample image from A
            imageB : sample image from B
        """
        self.steps += 1
        # Move the tensors to GPU
        imageA = imageA.to(self.device)
        imageB = imageB.to(self.device)

        # Do forward propagation through the generator network
        fakeB = self.generator_B(imageA)
        fakeA = self.generator_A(imageB)
        reconstructedB = self.generator_B(fakeA)
        reconstructedA = self.generator_A(fakeB)

        # Cycle loss
        forward_cycle_loss = self.cycle_loss(reconstructedA, imageA) * 10
        reverse_cycle_loss = self.cycle_loss(reconstructedB, imageB) * 10

        # Labels
        real_label = 1.0
        fake_label = 0.0

        # Generator loss
        target_real = torch.tensor(real_label)
        discriminator_A_preds = self.discriminator_A(fakeA) # Get discriminator predictions for fake images generated for A
        discriminator_B_preds = self.discriminator_B(fakeB) # Get discriminator predictions for fake images generated for B
        target_real = target_real.expand_as(discriminator_A_preds).to(self.device) # Target for generator loss is all 1's

        generator_A_loss = self.gan_loss(discriminator_A_preds, target_real)
        generator_B_loss = self.gan_loss(discriminator_B_preds, target_real)

        # Total generator loss
        generator_loss = generator_A_loss + generator_B_loss + forward_cycle_loss + reverse_cycle_loss

        # Discriminator loss
        fakeB = self.generator_B_buffer.add_and_get_image(fakeB).detach() # Get images from buffer for generator B
        fakeA = self.generator_A_buffer.add_and_get_image(fakeA).detach() # Get images from buffer for generator A
        
        pred_real_A = self.discriminator_A(imageA)
        pred_fake_A = self.discriminator_A(fakeA)

        target_real = torch.tensor(real_label)
        target_fake = torch.tensor(fake_label)
        target_real = target_real.expand_as(pred_real_A).to(self.device)
        target_fake = target_fake.expand_as(pred_fake_A).to(self.device)

        discriminator_A_loss = self.gan_loss(pred_real_A, target_real) + self.gan_loss(pred_fake_A, target_fake)

        pred_real_B = self.discriminator_B(imageB)
        pred_fake_B = self.discriminator_B(fakeB)

        discriminator_B_loss = self.gan_loss(pred_real_B, target_real) + self.gan_loss(pred_fake_B, target_fake)

        # Update generators
        self.set_grads([self.discriminator_A, self.discriminator_B], False)
        self.generator_A_optimizer.zero_grad()
        self.generator_B_optimizer.zero_grad()  
        generator_loss.backward()
        self.generator_A_optimizer.step()
        self.generator_B_optimizer.step()

        # Update discriminators    
        self.set_grads([self.discriminator_A, self.discriminator_B], True)
        self.discriminator_A_optimizer.zero_grad()
        self.discriminator_B_optimizer.zero_grad() 
        discriminator_A_loss.backward()
        discriminator_B_loss.backward()
        self.discriminator_A_optimizer.step()
        self.discriminator_B_optimizer.step()

        print("Epoch : %d, Iteration : %d, Generator A loss : %0.3f, Generator B loss : %0.3f,  Discriminator A loss : %0.3f, Discriminator B loss : %0.3f, Forward cycle loss : %0.3f, Reverse cycle loss : %0.3f"
                    % (epoch, iteration, generator_A_loss.item(), generator_B_loss.item(), discriminator_A_loss.item(), discriminator_B_loss.item(), forward_cycle_loss.item(), reverse_cycle_loss.item()))

        # Write to tensorboard
        self.writer.add_scalar('Generator A loss', generator_A_loss.item(), self.steps)
        self.writer.add_scalar('Generator B loss', generator_B_loss.item(), self.steps)
        self.writer.add_scalar('Discriminator A loss', discriminator_A_loss.item(), self.steps)
        self.writer.add_scalar('Discriminator B loss', discriminator_B_loss.item(), self.steps)
        self.writer.add_scalar('Forward cycle loss', forward_cycle_loss.item(), self.steps)
        self.writer.add_scalar('Reverse cycle loss', reverse_cycle_loss.item(), self.steps)

        # Write images to tensorboard
    
    def scheduler_step(self):
        """ Update learning rate for all optimizers according to the scheduler
        """
        self.generator_A_scheduler.step()
        self.generator_B_scheduler.step()
        self.discriminator_A_scheduler.step()
        self.discriminator_B_scheduler.step()
    
    def train(self):
        for epoch in range(201):
            for i, data in enumerate(self.dataloader):
                imageA = data['A']
                imageB = data['B']

                self.optimize(imageA, imageB, epoch, i)    
            # Update learning rate at the end of the epoch
            self.scheduler_step()

            imagesA, imagesB = self.dataset.load_test_images() # Load test images
            
            imagesA = imagesA.to(self.device)
            imagesB = imagesB.to(self.device)
            # Generate fake images
            fakesB = self.generator_B(imagesA)
            fakesA = self.generator_A(imagesB)
            self.writer.add_image("Fake images A", vutils.make_grid(fakesA.data[:5]), epoch)
            self.writer.add_image("Real images A", vutils.make_grid(imagesA.data[:5]), epoch)
            self.writer.add_image("Fake images B", vutils.make_grid(fakesB.data[:5]), epoch)
            self.writer.add_image("Real images B", vutils.make_grid(imagesB.data[:5]), epoch)


            if epoch % 25 == 0:
                self.save_network(epoch)

    def load_networks(self, epoch):
        path = os.path.join(os.getcwd(), 'models', 'cyclegan')

        self.generator_A.load_state_dict(torch.load('%s/generator_A_epoch_%d.pth' % (path, epoch)))
        self.generator_B.load_state_dict(torch.load('%s/generator_B_epoch_%d.pth' % (path, epoch)))
        self.discriminator_A.load_state_dict(torch.load('%s/discriminator_A_epoch_%d.pth' % (path, epoch)))
        self.discriminator_B.load_state_dict(torch.load('%s/discriminator_B_epoch_%d.pth' % (path, epoch)))

    def test(self, folder):

        self.load_networks(epoch = 200) # Load the saved weights for all networks

        # Switch generator networks to eval mode for testing
        self.generator_A.eval()
        self.generator_B.eval()

        # Get the dataset and dataloaders
        self.dataset = CycleGANTestDataset(base_dir = '/home/paurvi/CycleGAN/datasets/summer2winter_yosemite', folder = folder)
        self.dataloader = DataLoader(self.dataset, batch_size = 1, num_workers = 2)

        for i, data in enumerate(self.dataloader):
            print(i)
            image = data['data']
            image = image.to(self.device)   
            if folder == 'A':
                fake = self.generator_B(image)
            else:
                fake = self.generator_A(image)
            path = data['path']
            print(path)
            file_name = path[0][59:]

            path = os.path.join(os.getcwd(), 'results', 'cyclegan')
            print('sddddddd',path)
            transform = torchvision.transforms.ToPILImage()
            fake = fake[0, :, :, :]
            image = image[0, :, :, :]
            fake = fake.cpu()

            fake = transform(fake)
            image = image.cpu()
            image = transform(image)
            if folder == 'A':
                real_path = os.path.join(path, 'realA')
                fake_path = os.path.join(path, 'fakeB')
            else:
                real_path = os.path.join(path, 'realB')
                fake_path = os.path.join(path, 'fakeA')

            real_path = real_path + file_name
            fake_path = fake_path + file_name

            fake.save(fake_path)
            image.save(real_path)

if __name__ == "__main__":
    cyclegan = GAN()
    cyclegan.test('B')
