#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Apr 30 23:26:02 2021

@author: valentinahrtonova
"""
from numpy import load, expand_dims
from numpy.random import randint
from matplotlib import pyplot
import os
        
def load_dataset(self, dataset_path):
    """Load and preprocess NPZ dataset"""
             
    data = load(dataset_path)
    img_tag, img_A, img_B = data['arr_0'], data['arr_1'], data['arr_2']

    print(img_A.dtype, img_B.dtype)
    
    # Scale images from [0,1] to [-1,1]
    img_A = (img_A - 0.5) * 2 
    img_B = (img_B - 0.5) * 2 
    
    img_A = expand_dims(img_A, axis=3)
    img_B = expand_dims(img_B, axis=3)
    
    return [img_tag, img_A, img_B]

def sample_images(self, epoch, batch):
    """Use generator to sample images"""
    
    os.makedirs(self.directory , exist_ok=True)
    
    img_tag, img_A, img_B = self.dataset_test
    
    index = randint(0, img_A.shape[0], 3)
    imgs_A, imgs_B = img_A[index], img_B[index]
    
    # Generate images based on imgs_A
    fake_B = self.generator.predict(imgs_A)   
    
    # Scale from [-1,1] to [0,1]
    real_A = (imgs_A + 1) / 2.0
    real_B = (imgs_B + 1) / 2.0
    fake_B = (fake_B + 1) / 2.0
    
    n_samples = 3
        
    pyplot.figure(figsize=(20, 10))
    pyplot.style.use('grayscale')
    
    # Plot masks
    for i in range(n_samples):
        pyplot.subplot(3, n_samples, 1 + i)
        pyplot.title('Mask')
        pyplot.axis('off')
        pyplot.imshow(real_A[i])
	# Plot generated images
    for i in range(n_samples):
        pyplot.subplot(3, n_samples, 1 + n_samples + i)
        pyplot.title('Generated')
        pyplot.axis('off')
        pyplot.imshow(fake_B[i])
	# Plot ground truth images
    for i in range(n_samples):
        pyplot.subplot(3, n_samples, 1 + n_samples*2 + i)
        pyplot.title('Ground truth')
        pyplot.axis('off')
        pyplot.imshow(real_B[i])
        
	# save plot to file
    filename1 = self.directory + '/plot_%06d.png' % (epoch)
    pyplot.savefig(filename1)
    
	# save models
    g_filename = self.directory + '/g_model_%03d.h5' % (epoch)
    self.generator.save(g_filename)
    d_filename = self.directory + '/d_model_%03d.h5' % (epoch)
    self.discriminator.save(d_filename)
    gan_filename = self.directory + '/gan_model_%03d.h5' % (epoch)
    self.combined.save(gan_filename)
    print('>Saved: %s' % (g_filename))

def plot_losses(self):
	# plot history
    pyplot.figure(figsize=(10,5))
    pyplot.plot(self.d_loss_real, label='d_loss_real')
    pyplot.plot(self.d_loss_fake, label='d_loss_fake')
    pyplot.plot(self.d_loss, label='d_loss')
    pyplot.plot(self.g_loss, label='g_loss')
    pyplot.title("Training loss")
    pyplot.xlabel('Epoch')
    pyplot.ylabel('Loss')
    pyplot.legend()
    pyplot.show()
    pyplot.savefig(self.directory + '/plot_loss.png')
    pyplot.close()