import numpy as np
import itertools
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision.utils import save_image
from logger import Logger
LOGGER = Logger().logger()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
from discriminator import Discriminator
from generator import Generator
from loss import Loss
from preprocess import LoadData
from tqdm import tqdm
from utils import initialize_weights
LOGGER.info("Cuda status: "+str(device))
torch.cuda.empty_cache()
import os
class Train():
    def __init__(self,data="dataset/summer2winter_yosemite",pair=False,epochs=200,batch_size=1):
        '''
        @params
        @return
        '''
        if not os.path.exists("images"):
            os.mkdir("images")
        if not os.path.exists("weights"):
            os.mkdir("weights")
        self.epochs = epochs
        self.batch_size = batch_size

        self.gen_XY = Generator().to(device)
        self.gen_YX = Generator().to(device)
        self.dis_X = Discriminator().to(device)
        self.dis_Y = Discriminator().to(device)
        self.gen_XY.apply(initialize_weights)
        self.gen_YX.apply(initialize_weights)
        self.dis_X.apply(initialize_weights)
        self.dis_Y.apply(initialize_weights)

        self.gen_XY_optim = torch.optim.Adam(itertools.chain(self.gen_XY.parameters(),self.gen_YX.parameters()), lr=0.0002, betas=(0.5, 0.999))
        self.dis_X_optim = torch.optim.Adam(self.dis_X.parameters(), lr=0.0002, betas=(0.5, 0.999))
        self.dis_Y_optim = torch.optim.Adam(self.dis_Y.parameters(), lr=0.0002, betas=(0.5, 0.999))

        self.gen_XY_scheduler = torch.optim.lr_scheduler.StepLR(self.gen_XY_optim, step_size=100, gamma=0.1)

        self.dis_X_scheduler = torch.optim.lr_scheduler.StepLR(self.dis_X_optim, step_size=100, gamma=0.1)
        self.dis_Y_scheduler = torch.optim.lr_scheduler.StepLR(self.dis_Y_optim, step_size=100, gamma=0.1)

        self.cycle_loss = Loss().cycle_consistency
        self.identity_loss = Loss().identity
        self.adversarial_loss = Loss().adversarial_G

        self.losses = {"G": [], "D": [], "C": [], "I": [], "T": []}
        self.dataset = LoadData(data=data,pair=pair)
        self.dataloader = DataLoader(self.dataset, batch_size=self.batch_size, shuffle=True, num_workers=4)


    def train(self):
        '''
        @params
        @return
        '''
        EPOCHS = self.epochs
        batch_size = self.batch_size
        for epoch in range(EPOCHS):
            '''
            Steps:
            For Generator
            1. Load Images
            2. Reset Gradients
            3. Generate Fake Images
            4. Calculate Advetisarial Loss
            5. Generate Original Images from Fake Images
            6. Calculate Cycle Consistency Loss
            7. Total Loss
            8. Backpropagate
            9. Update Weights

            For Discriminator
            10. Reset Gradients
            11. Calculate Adversarial Loss between labels for real images and generated images
            12. Total loss = Average of both losses
            13. Backpropagate
            14. Update Weights
            15. Do same for other discriminator

            16. Update Logs
            17. Save Image
            18. Update Learning Rate

            '''
            adversarial_loss = self.adversarial_loss()
            cycle_loss = self.cycle_loss()
            identity_loss = self.identity_loss()
            size = len(self.dataloader)
            for i, data in tqdm(enumerate(self.dataloader), total=size):
                torch.cuda.empty_cache()
                real_X = data['X'].to(device)
                real_Y = data['Y'].to(device)
                #print(real_X.shape)
                #print(real_Y.shape)
                batch_size = real_X.size(0)
                real_label = torch.ones(batch_size, 1).to(device)
                fake_label = torch.zeros(batch_size, 1).to(device)
                #print(real_label.shape)

                # Training the generator
                self.gen_XY_optim.zero_grad()

                identity_X = self.gen_YX(real_X)
                identity_Y = self.gen_XY(real_Y)

                loss_iden_X = identity_loss(identity_X,real_X) * 10
                loss_iden_Y = identity_loss(identity_Y,real_Y) * 10


                fake_gen_X = self.gen_YX(real_Y)
                fake_gen_Y = self.gen_XY(real_X)

                fake_gen_X_label = self.dis_X(fake_gen_X)
                fake_gen_Y_label = self.dis_Y(fake_gen_Y)
                #print(fake_gen_X.shape,fake_gen_X_label.shape)
                #print(fake_gen_Y_label,real_label)
                loss_gen_Y2X = adversarial_loss(fake_gen_X_label, real_label)
                loss_gen_X2Y = adversarial_loss(fake_gen_Y_label, real_label)

                recovered_Y = self.gen_XY(fake_gen_X)
                recovered_X = self.gen_YX(fake_gen_Y)

                #print(recovered_Y.shape,recovered_X.shape)

                loss_cycle_Y2X = cycle_loss(recovered_Y, real_Y) * 20
                loss_cycle_X2Y = cycle_loss(recovered_X, real_X) * 20

                total_loss = loss_gen_Y2X + loss_gen_X2Y + loss_cycle_Y2X + loss_cycle_X2Y + loss_iden_X + loss_iden_Y
                #backprop
                total_loss.backward()
                self.gen_XY_optim.step()

                # Training the discriminator

                # Discriminator for X
                self.dis_X_optim.zero_grad()

                real_X_label = self.dis_X(real_X)
                loss_dis_real_X = adversarial_loss(real_X_label, real_label)

                fake_X_label = self.dis_X(fake_gen_X.detach())
                loss_dis_fake_X = adversarial_loss(fake_X_label, fake_label)

                loss_dis_X = (loss_dis_real_X + loss_dis_fake_X) / 2
                #backprop
                loss_dis_X.backward()
                self.dis_X_optim.step()

                # Discriminator for Y
                self.dis_Y_optim.zero_grad()

                real_Y_label = self.dis_Y(real_Y)
                loss_dis_real_Y = adversarial_loss(real_Y_label, real_label)

                fake_Y_label = self.dis_Y(fake_gen_Y.detach())
                loss_dis_fake_Y = adversarial_loss(fake_Y_label, fake_label)

                loss_dis_Y = (loss_dis_real_Y + loss_dis_fake_Y) / 2

                #backprop
                loss_dis_Y.backward()
                self.dis_Y_optim.step()

                # Update Logs
                self.losses["G"].append(total_loss.item())
                self.losses["D"].append((loss_dis_X.item() + loss_dis_Y.item()) / 2)
                self.losses["C"].append((loss_cycle_Y2X.item() + loss_cycle_X2Y.item()) / 2)
                LOGGER.info("Epoch: {} | i: {} | G: {} | D: {} | C: {}".format(epoch, i, total_loss.item(), (loss_dis_X.item() + loss_dis_Y.item()) / 2, (loss_cycle_Y2X.item() + loss_cycle_X2Y.item()) / 2))

                # Save Image
                if i % 100 == 0:
                    save_image(fake_gen_X, "images/{}_{}_fake_X.png".format(epoch, i))
                    save_image(fake_gen_Y, "images/{}_{}_fake_Y.png".format(epoch, i))
                    save_image(real_X, "images/{}_{}_real_X.png".format(epoch, i))
                    save_image(real_Y, "images/{}_{}_real_Y.png".format(epoch, i))

            # Update Learning Rate
            self.gen_XY_scheduler.step()
            self.dis_X_scheduler.step()
            self.dis_Y_scheduler.step()

            # Save weights
            torch.save(self.gen_XY.state_dict(), "weights/gen_XY.pth")
            torch.save(self.gen_YX.state_dict(), "weights/gen_YX.pth")
            torch.save(self.dis_X.state_dict(), "weights/dis_X.pth")
            torch.save(self.dis_Y.state_dict(), "weights/dis_Y.pth")

        #Save losses
        torch.save(self.losses,"losses.pt")

if __name__ == "__main__":
    train = Train()
    train.train()
