#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue May 11 12:22:15 2021

@author: valentinahrtonova
"""

from numpy import load, expand_dims
from numpy.random import randint
from matplotlib import pyplot as plt
import os

from keras.models import load_model

def load_dataset(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']
    
    # 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 window_image(image, window_center, window_width):
    """Window image for CT range"""
    
    img_min = window_center - (window_width / 2)
    img_max = window_center + (window_width / 2)
    win_image = image.copy()
    win_image[win_image < img_min] = img_min
    win_image[win_image > img_max] = img_max

    return win_image

def reverse_normalize_image(image, img_min, img_max):
    """Reverse image normalization"""
    
    image = (image + img_min) * (img_max - img_min)
    renorm_image = image - 1024

    return renorm_image

"""TESTING DEMO - SAGITAL"""

directory = 'demo_output'

model = load_model('model_sagital_200.h5')
tag, real_A, real_B = load_dataset('pix2pix_sagital_input_demo.npz')

print('Loaded', real_A.shape, real_B.shape)

# generate a batch of fake samples
fake_B = model.predict(real_A)

# scale all pixels from [-1,1] to [0,1]
real_A = (real_A + 1) / 2.0
real_B = (real_B + 1) / 2.0
fake_B = (fake_B + 1) / 2.0

real_A = real_A[0,:,:,:]
real_B = real_B[0,:,:,:]
fake_B = fake_B[0,:,:,:]

renorm_fake = reverse_normalize_image(fake_B, 0, 3000)
bone_fake = window_image(renorm_fake, 400, 1000)
st_fake = window_image(renorm_fake, 50, 250)

fig, (axs) = plt.subplots(2, 3, constrained_layout=True, sharey=True)
#fig.suptitle(tag, fontsize=12)  
plt.style.use('grayscale')

ax1 = axs[0,0]
ax1.imshow(fake_B)
ax1.set_title('Generated')
ax1.axis('off')

ax2 = axs[0,1]
ax2.imshow(bone_fake)
ax2.set_title('Bone')
ax2.axis('off')

ax3 = axs[0,2]
ax3.imshow(st_fake)
ax3.set_title('Soft tissue')
ax3.axis('off')

renorm_real = reverse_normalize_image(real_B, 0, 3000)
bone_real = window_image(renorm_real, 400, 1000)
st_real = window_image(renorm_real, 50, 250)

ax4 = axs[1,0]
ax4.imshow(real_B)
ax4.set_title('Grount truth')
ax4.axis('off')

ax5 = axs[1,1]
ax5.imshow(bone_real)
ax5.set_title('Bone')
ax5.axis('off')

ax6 = axs[1,2]
ax6.imshow(st_real)
ax6.set_title('Soft tissue')
ax6.axis('off')
plt.close()

# save plot to file
filename1 = directory + '/plot_demo'
fig.savefig(filename1, transparent=True)

plt.imshow(real_A, cmap = 'gray')
plt.axis('off')
filename2 = directory + '/generated_full'
plt.savefig(filename2, transparent=True)

plt.imshow(bone_fake, cmap = 'gray')
plt.axis('off')
filename3 = directory + '/generated_bone'
plt.savefig(filename3, transparent=True)

plt.imshow(st_fake, cmap = 'gray')
plt.axis('off')
filename4 = directory + '/generated_soft_tissue'
plt.savefig(filename4, transparent=True)