#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Apr 30 20:45:43 2021

@author: valentinahrtonova
"""

from __future__ import print_function, division

import datetime
import numpy as np
from numpy.random import randint

from pix2pix_melanoma_models import build_gan
from pix2pix_melanoma_data_loader import load_dataset, sample_images

class Pix2Pix():
    def __init__(self):
        
        #Options
        self.epochs = 1000
        self.batch_size = 4
        self.sample_interval= 20
        
        self.directory = 'melanoma_May_19'
        path_dataset_train = 'melanoma_pix2pix.npz'
        path_dataset_test = 'melanoma_pix2pix_test.npz'
        
        # Load training and testing dataset
        self.dataset_train = load_dataset(self, path_dataset_train)
        print('Loaded training dataset: ', self.dataset_train[0].shape, self.dataset_train[1].shape)
        self.dataset_test = load_dataset(self, path_dataset_test)
        print('Loaded testing dataset:', self.dataset_train[0].shape, self.dataset_train[1].shape)
        
        # Input image shape
        self.img_shape = self.dataset_train[0].shape[1:]
        self.img_rows = self.img_shape[0]
        self.img_cols = self.img_shape[1]
        self.channels = self.img_shape[2]
        print ('Input image shape is: ', self.img_shape)
        
        # Build GAN model
        self.gan = build_gan(self)
        
    def train(self):
        """Train the GAN model"""

        start_time = datetime.datetime.now()
        
        epochs = self.epochs
        batch_size = self.batch_size
        sample_interval = self.sample_interval
        
        img_A, img_B = self.dataset_train

        # Adversarial loss ground truths
        gt_valid = np.ones((self.batch_size,) + self.disc_patch)
        gt_fake = np.zeros((self.batch_size,) + self.disc_patch)
        
        n_batches = int(len(self.dataset_train[0])/self.batch_size)
        
        print('>Epochs: %s, Batch per epoch: %s' % (self.epochs, n_batches))

        for epoch in range(self.epochs):
            
            for i in range(0,n_batches):
                
                index = randint(0, img_A.shape[0], batch_size)
                imgs_A, imgs_B = img_A[index], img_B[index]
                
                """Train Discriminator"""

                # Generate images based on imgs_A
                fake_B = self.generator.predict(imgs_A)

                # Train the discriminators
                d_loss_real = self.discriminator.train_on_batch([imgs_A, imgs_B], gt_valid)
                d_loss_fake = self.discriminator.train_on_batch([imgs_A, fake_B], gt_fake)
                d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)

                """Train Generator"""

                # Train the generators
                g_loss = self.combined.train_on_batch([imgs_A, imgs_B], [gt_valid, imgs_B])

                elapsed_time = datetime.datetime.now() - start_time
                
                # Plot the progress
                print ("[Epoch %d/%d][D loss: %f] [G loss: %f] time: %s" % (epoch, epochs,
                                                                        d_loss,
                                                                        g_loss[0],
                                                                        elapsed_time))

                # If at save interval => save generated image samples
                if i % sample_interval == 0:
                    sample_images(self, epoch, i)

if __name__ == '__main__':
    gan = Pix2Pix()
    gan.train()

