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

@author: valentinahrtonova
"""

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

from pix2pix_axial_models import build_gan
from pix2pix_axial_data_loader import load_dataset, sample_images, plot_losses

class Pix2Pix():
    def __init__(self):
        
        #Training options
        self.epochs = 1000
        self.batch_size = 4
        self.sample_interval = 20
        
        self.directory = 'spine_axial_May_14'
        path_dataset_train = 'verse_dataset_axial_final_train.npz'
        path_dataset_test = 'verse_dataset_axial_final_test.npz'
        
        # Load training and testing dataset
        self.dataset_train = load_dataset(self, path_dataset_train)
        self.dataset_train_size = self.dataset_train[1].shape[0]
        print('Loaded training dataset: ', self.dataset_train[1].shape, self.dataset_train[2].shape)
        self.dataset_test = load_dataset(self, path_dataset_test)
        self.dataset_test_size = self.dataset_test[1].shape[0]
        print('Loaded testing dataset:', self.dataset_train[1].shape, self.dataset_train[2].shape)
        
        # Input image shape
        self.img_shape = self.dataset_train[1].shape[1:]
        print ('Input image shape is: ', self.img_shape)
        self.img_rows = self.img_shape[0]
        self.img_cols = self.img_shape[1]
        self.channels = self.img_shape[2]
        
        # 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_tag, img_A, img_B = self.dataset_train

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

        for epoch in range(epochs):
            
            for i in range(0, n_batches):
                
                # sample batch
                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_real)
                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_real, imgs_B])

                elapsed_time = datetime.datetime.now() - start_time

                # Save losses
                self.d_loss_real.append(d_loss_real)
                self.d_loss_fake.append(d_loss_fake)
                self.d_loss.append(d_loss)
                self.g_loss.append(g_loss)

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

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

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

