#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed May 26 16:23:50 2021

@author: valentinahrtonova
"""

from numpy import load, flipud, fliplr, zeros, ones
from numpy.random import randn, randint, random_sample

from matplotlib import pyplot
import os
 
def augment_data(X_real):
    for i in range(X_real.shape[0]):
        if random_sample(1) > 0.2:
            X_real[i] = flipud(X_real[i])

        if random_sample(1) > 0.2:
            X_real[i] = fliplr(X_real[i])
            
        return X_real
    
def load_dataset(self, dataset_path):
    """Load and preprocess NPZ dataset"""
     
    # Load compressed arrays             
    data = load(dataset_path)
    imgs = data['arr_0']
    imgs.astype('float32')
    
    # Scale images from [0,255] to [-1,1]
    imgs = (imgs - 127.5) / 127.5
    
    return imgs

# select real samples
def generate_real_samples(self, n_samples):
	# choose random instances
	ix = randint(0, self.dataset.shape[0], n_samples)
	# retrieve selected images
	X = self.dataset[ix]
	# generate 'real' class labels (1)
	y = ones((n_samples, 1))
	return X, y

# generate points in latent space as input for the generator
def generate_latent_points(self, n_samples):
	# generate points in the latent space
	x_input = randn(self.latent_dim * n_samples)
	# reshape into a batch of inputs for the network
	x_input = x_input.reshape(n_samples, self.latent_dim)
	return x_input

# use the generator to generate n fake examples, with class labels
def generate_fake_samples(self, n_samples):
	# generate points in latent space
	x_input = generate_latent_points(self, n_samples)
	# predict outputs
	X = self.generator.predict(x_input)
	# create 'fake' class labels (0)
	y = zeros((n_samples, 1))
	return X, y

# evaluate the discriminator, plot generated images, save generator model
def save_progress(self, epoch, n_samples=30):
	# prepare real samples
	X_real, y_real = generate_real_samples(self, n_samples)
    # augment real samples
	X_real = augment_data(X_real)
	# prepare fake examples
	X_fake, y_fake = generate_fake_samples(self, n_samples)

	# save plot
	os.makedirs(self.directory, exist_ok=True)
	# scale from [-1,1] to [0,1]
	examples = (X_fake + 1) / 2.0
    
	n = 5
    
	# plot images
	for i in range(n * n):
		# define subplot
		pyplot.subplot(n, n, 1 + i)
		# turn off axis
		pyplot.axis('off')
		# plot raw pixel data
		pyplot.imshow(examples[i])
	# save plot to file
	#pyplot.show()
	filename = self.directory + '/generated_e%03d.png' % (epoch+1)
	pyplot.savefig(filename)
	print('>Saved: %s' % (filename))
	pyplot.close()
    
	# save the generator model tile file
	g_filename = self.directory + '/g_model_%03d.h5' % (epoch+1)
	self.generator.save(g_filename)
	d_filename = self.directory + '/d_model_%03d.h5' % (epoch+1)
	self.discriminator.save(d_filename)
	gan_filename = self.directory + '/gan_model_%03d.h5' % (epoch+1)
	print('>Saved: %s' % (g_filename))
	self.combined.save(gan_filename)