#!/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"""
     
    # Load compressed arrays             
    data = load(dataset_path)
    img_A, img_B = data['arr_0'], data['arr_1']
    
    # Covert to float
    img_A.astype('float32')
    img_B.astype('float32')
    
    # Scale images from [0,255] to [-1,1]
    img_A = (img_A - 127.5) / 127.5
    img_B = (img_B - 127.5) / 127.5
    
    return [img_A, img_B]


def sample_images(self, epoch, batch):
    """Use generator to sample images"""
    
    os.makedirs(self.directory , exist_ok=True)
    
    dataset = self.dataset_test
    img_A, img_B = dataset
    
    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 all pixels 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
    
    # plot real source images
    for i in range(n_samples):
        pyplot.subplot(3, n_samples, 1 + i)
        pyplot.axis('off')
        pyplot.imshow(real_A[i])
	# plot generated target image
    for i in range(n_samples):
        pyplot.subplot(3, n_samples, 1 + n_samples + i)
        pyplot.axis('off')
        pyplot.imshow(fake_B[i])
	# plot real target image
    for i in range(n_samples):
        pyplot.subplot(3, n_samples, 1 + n_samples*2 + i)
        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))
