#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Nov 14 21:47:22 2018

@author: Josefine
"""

import tensorflow as tf
import numpy as np
import glob
import re
from skimage.transform import resize

imgDim = 256
labelDim = 256

##############################################################################
###                              Data functions                         ######
##############################################################################
def natural_sort(l): 
    convert = lambda text: int(text) if text.isdigit() else text.lower() 
    alphanum_key = lambda key: [ convert(c) for c in re.split('([0-9]+)', key) ] 
    return sorted(l, key = alphanum_key)

def create_data(filename_img,direction):
    images = []
    file = np.load(filename_img)
    a = file['images']
    # Normalize:
    #a2 = np.clip(a,-1000,1000)
    #a3 = np.interp(a2, (a2.min(), a2.max()), (-1, +1))
    im = resize(a,(labelDim,labelDim,labelDim),order=0)
    if direction == 'axial':
        for i in range(im.shape[0]):
            images.append((im[i,:,:]))
    if direction == 'sag':
        for i in range(im.shape[1]):
            images.append((im[:,i,:]))
    if direction == 'cor':
        for i in range(im.shape[2]):
            images.append((im[:,:,i]))    
    images = np.asarray(images)
    images = images.reshape(-1, imgDim,imgDim,1)
    return images

# Load test data
filelist_test = natural_sort(glob.glob('WHS/Data/test_segments_*.npz')) # list of file names

#############################################################################
##                  Reload network and predict                         ######
#############################################################################
#
## =============================================================================
print("====================== LOAD AXIAL NETWORK: ===========================")
# Doing predictions with the model 
tf.reset_default_graph()      

new_saver = tf.train.import_meta_graph('WHS/Results/segmentation/model_axial/model.ckpt.meta')
 
prediction = np.zeros([1,256,256,9])
with tf.Session() as sess:
    new_saver.restore(sess, tf.train.latest_checkpoint('WHS/Results/segmentation/model_axial/'))
    graph = tf.get_default_graph()       
    x = graph.get_tensor_by_name("x_train:0")
    op_to_restore = graph.get_tensor_by_name("output/Softmax:0")
    keep_rate = graph.get_tensor_by_name("Placeholder:0")
    context = graph.get_tensor_by_name("concat_5:0")
    x_contextual = graph.get_tensor_by_name("x_train_context:0")
    for i in range(30,len(filelist_test)):
        print('Processing test image', (i+1),'out of',(np.max(range(len(filelist_test)))+1))
        # Find renderings corresponding to the given name
        prob_maps = []
        x_test = create_data(filelist_test[i],'axial')
        for k in range(x_test.shape[0]):
            x_test_image = np.expand_dims(x_test[k,:,:,:], axis=0)
            y_output,out_context = sess.run([tf.nn.softmax(op_to_restore),context], feed_dict={x: x_test_image, x_contextual: prediction,keep_rate: 1.0})
            prediction[0,:,:,:] = out_context
            prob_maps.append(y_output[0,:,:,:])
        np.savez('WHS/Results/Predictions/segment/train_prob_maps_axial_{}'.format(i),prob_maps=prob_maps)                            
print("================ DONE WITH AXIAL PREDICTIONS! ==================")  
#
# =============================================================================
#print("====================== LOAD SAGITTAL NETWORK: ===========================")
## Doing predictions with the model 
#tf.reset_default_graph()      
#
#new_saver = tf.train.import_meta_graph('WHS/Results/segmentation/model_sag/model.ckpt.meta')
#prediction = np.zeros([1,256,256,9])
#with tf.Session() as sess:
#    new_saver.restore(sess, tf.train.latest_checkpoint('WHS/Results/segmentation/model_sag/'))
#    graph = tf.get_default_graph()       
#    x = graph.get_tensor_by_name("x_train:0")
#    keep_rate = graph.get_tensor_by_name("Placeholder:0")
#    op_to_restore = graph.get_tensor_by_name("output/Softmax:0")
#    context = graph.get_tensor_by_name("concat_5:0")
#    x_contextual = graph.get_tensor_by_name("x_train_context:0")
#    for i in range(30,len(filelist_test)):
#        print('Processing test image', (i+1),'out of',(np.max(range(len(filelist_test)))+1))
#        # Find renderings corresponding to the given name
#        prob_maps = []
#        x_test = create_data(filelist_test[i],'sag')
#        for k in range(x_test.shape[0]):
#            x_test_image = np.expand_dims(x_test[k,:,:,:], axis=0)
#            y_output,out_context = sess.run([tf.nn.softmax(op_to_restore),context], feed_dict={x: x_test_image, x_contextual: prediction,keep_rate: 1.0})
#            prediction[0,:,:,:] = out_context
#            prob_maps.append(y_output[0,:,:,:])
#        np.savez('WHS/Results/Predictions/segment/train_prob_maps_sag_{}'.format(i),prob_maps=prob_maps)                            
#print("================ DONE WITH SAGITTAL PREDICTIONS! ==================")  
##
#print("====================== LOAD CORONAL NETWORK: ===========================")
## Doing predictions with the model 
#tf.reset_default_graph()      
#
#new_saver = tf.train.import_meta_graph('WHS/Results/segmentation/model_cor/model.ckpt.meta')
#prediction = np.zeros([1,256,256,9])
#with tf.Session() as sess:
#    new_saver.restore(sess, tf.train.latest_checkpoint('WHS/Results/segmentation/model_cor/'))
#    graph = tf.get_default_graph()       
#    x = graph.get_tensor_by_name("x_train:0")
#    keep_rate = graph.get_tensor_by_name("Placeholder:0")
#    op_to_restore = graph.get_tensor_by_name("output/Softmax:0")
#    context = graph.get_tensor_by_name("concat_5:0")
#    x_contextual = graph.get_tensor_by_name("x_train_context:0")
#    for i in range(30,len(filelist_test)):
#        print('Processing test image', (i+1),'out of',(np.max(range(len(filelist_test)))+1))
#        # Find renderings corresponding to the given name
#        prob_maps = []
#        x_test = create_data(filelist_test[i],'cor')
#        for k in range(x_test.shape[0]):
#            x_test_image = np.expand_dims(x_test[k,:,:,:], axis=0)
#            y_output,out_context = sess.run([tf.nn.softmax(op_to_restore),context], feed_dict={x: x_test_image, x_contextual: prediction,keep_rate: 1.0})
#            prediction[0,:,:,:] = out_context
#            prob_maps.append(y_output[0,:,:,:])
#        np.savez('WHS/Results/Predictions/segment/train_prob_maps_cor_{}'.format(i),prob_maps=prob_maps)                            
#print("================ DONE WITH CORONAL PREDICTONS! ==================")  
#