#!/usr/bin/env python3

# -*- coding: utf-8 -*-

"""

Created on Tue Nov 20 19:44:39 2018

@author: Josefine

"""

import numpy as np

import re

import nibabel as nib

import glob

from skimage.transform import resize

from scipy import ndimage

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)

# Create original high res data function:

def create_data(filename_img,filename_label):

    images = []

    a = nib.load(filename_img)

    a = a.get_data()

    # Normalize:

    a2 = np.clip(a,-1000,1000)

    a3 = np.interp(a2, (a2.min(), a2.max()), (-1, +1))

    # Reshape:

    img = np.zeros([512,512,512])+np.min(a3)

    index1 = int(np.ceil((512-a.shape[2])/2))

    index2 = int(512-np.floor((512-a.shape[2])/2))

    img[:,:,index1:index2] = a3

    for i in range(img.shape[2]):

            images.append((img[:,:,i]))

    images = np.asarray(images)

    return images

# Fusion of low resolution probablity maps

def fusion(prob_maps_axial, prob_maps_sag, prob_maps_cor):

    sag_to_axial = []

    for i in range(prob_maps_sag.shape[2]):

        sag_to_axial.append((prob_maps_sag[:,i,:,:]))  

    sag_to_axial = np.asarray(sag_to_axial)

    # Reshape coronal data to match axial:

    cor_to_sag = []

    for i in range(prob_maps_cor.shape[2]):

        cor_to_sag.append((prob_maps_cor[:,i,:,:]))  

    cor_to_sag = np.asarray(cor_to_sag)

    cor_to_axial = []

    for i in range(prob_maps_cor.shape[2]):

        cor_to_axial.append((cor_to_sag[:,:,i,:]))  

    cor_to_axial = np.asarray(cor_to_axial)

    cor_to_axial2 = []

    for i in range(prob_maps_cor.shape[2]):

        cor_to_axial2.append((cor_to_axial[:,i,:,:]))  

    cor_to_axial = np.asarray(cor_to_axial2)

    temp = np.maximum.reduce([sag_to_axial,cor_to_axial,prob_maps_axial])

    return temp

def remove_objects(binary_mask):

    labelled_mask, num_labels = ndimage.label(binary_mask)

    # Let us now remove all the too small regions.

    refined_mask = binary_mask.copy()

    minimum_cc_sum = 5000

    for label in range(num_labels):

        if np.sum(refined_mask[labelled_mask == label]) < minimum_cc_sum:

            refined_mask[labelled_mask == label] = 0

    return refined_mask

filelist_train = natural_sort(glob.glob('WHS/ct_train_test/ct_test/*_image.nii.gz')) # list of file names

cropped_files = natural_sort(glob.glob('WHS/Data/test_segments_*.npz')) # list of file names

files_axial = natural_sort(glob.glob('WHS/Results/Predictions/segment/train_prob_maps_axial_*.npz')) # list of file names

files_sag = natural_sort(glob.glob('WHS/Results/Predictions/segment/train_prob_maps_sag_*.npz')) # list of file names

files_cor = natural_sort(glob.glob('WHS/Results/Predictions/segment/train_prob_maps_cor_*.npz')) # list of file names

for n in range(len(files_axial)):

    axial_data = np.load(files_axial[n])

    prob_maps_axial = axial_data['prob_maps']

    sag_data = np.load(files_sag[n])

    prob_maps_sag = sag_data['prob_maps']

    cor_data = np.load(files_cor[n])

    prob_maps_cor = cor_data['prob_maps']

    cut_file = np.load(cropped_files[n])

    cut = cut_file['cut']

    # Create fused propability map

    fused_prob_maps = fusion(prob_maps_axial, prob_maps_sag, prob_maps_cor)

    side_length = cut[1]-cut[0]

    lab = np.zeros([side_length,side_length,side_length,8])

    for i in range(8):

        lab[:,:,:,i] = resize(fused_prob_maps[:,:,:,i],(side_length,side_length,side_length))

    full_labels = np.zeros([512,512,512,8])

    full_labels[cut[0]:cut[1],cut[0]:cut[1],cut[0]:cut[1],:] = lab

    labels = full_labels.argmax(axis=-1)

    print('Test image', (n+1))

    np.savez('WHS/Results/Predictions/final/prediction_{}'.format(n),prob_map = labels)