
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Created on Tue May 11 22:53:23 2021

@author: valentinahrtonova
"""
from __future__ import print_function, division

from keras.initializers import RandomNormal
from keras.optimizers import Adam
from keras.losses import BinaryCrossentropy

from numpy import ones
import datetime

from DCGAN_models_256 import build_gan
from DCGAN_data_loader import augment_data, load_dataset, generate_real_samples, generate_latent_points, generate_fake_samples, save_progress

class DCGAN():
    def __init__(self):
        self.img_rows = 256
        self.img_cols = 256
        self.channels = 3
        self.img_shape = (self.img_rows, self.img_cols, self.channels)
        self.latent_dim = 100
        
        dataset_path = 'melanoma_pix2pix.npz' 
        
        self.directory = 'DCGAN_generated'
        self.dataset = load_dataset(self, dataset_path)

        self.opt = Adam(lr=0.0002, beta_1=0.5)
        self.init = RandomNormal(mean=0.0, stddev=0.02, seed=None)
        
        self.gf = 32
        self.df = 32

        # Combined GAN model
        self.combined = build_gan(self)
        self.combined.compile(loss=BinaryCrossentropy(label_smoothing = 0.1),
        optimizer=self.opt)

    def train(self, epochs, n_batch=16, sample_interval=20):

        bat_per_epo = int(self.dataset.shape[0] / n_batch)	
        start_time = datetime.datetime.now()
        
        for i in range(epochs):
            for j in range(bat_per_epo):
                
                ## Update discriminator
                X_real, y_real = generate_real_samples(self, n_batch)
                
                # Augment real samples
                X_real = augment_data(X_real)
    			
                d_loss1 = self.discriminator.train_on_batch(X_real, y_real)
                X_fake, y_fake = generate_fake_samples(self, n_batch)
                d_loss2 = self.discriminator.train_on_batch(X_fake, y_fake)
                
                ## Update the generator
    
    			# Prepare input for generator
                X_gan = generate_latent_points(self, n_batch)
    			# Prepare labels for fake samples
                y_gan = ones((n_batch, 1))
    			
                g_loss = self.combined.train_on_batch(X_gan, y_gan)
                
                elapsed_time = datetime.datetime.now() - start_time
                print('>%d, %d/%d, d1=%.3f, d2=%.3f g=%.3f time: %s'  %
    				(i+1, j+1, bat_per_epo, d_loss1, d_loss2, g_loss, elapsed_time))
                
    		# evaluate the model performance, sometimes
            if (i+1) % sample_interval == 0:
                save_progress(self, i)


if __name__ == '__main__':
    dcgan = DCGAN()
    dcgan.train(epochs=4000, n_batch=4, sample_interval=5)
