# -*- coding: utf-8 -*-
"""
Created on Thu Jun 25 10:01:12 2020

@author: t_estienne
"""
import os 
import SimpleITK as sitk
from tqdm import tqdm 

home = './data/'
main_path = home + 'Pancreas-CT/'
save_path = home + 'Pancreas-CT/resample/'


end = '.nii.gz'
n = len(end)

def load_sitk(path):
    img = sitk.ReadImage(path)
    return img

def load_patient(patient, folder, is_label=False):
    
    if is_label:
        path = folder + patient +'-seg' + end
    else:
        path = folder + patient + end
        
    ct = load_sitk(path)

    return ct

def do_resampling(image, spacing_x, spacing_y, spacing_z, interpolator='Linear'):

    if interpolator == 'Linear':
        interpolator = sitk.sitkLinear
    else:
        interpolator = sitk.sitkNearestNeighbor

    # to calcul new dimensions
    spacing = image.GetSpacing()
    size = image.GetSize()
    fact_x = spacing[0] / spacing_x
    fact_y = spacing[1] / spacing_y
    fact_z = spacing[2] / spacing_z
    size_x = int(round(size[0] * fact_x))
    size_y = int(round(size[1] * fact_y))
    size_z = int(round(size[2] * fact_z))
    # to do resampling
    f = sitk.ResampleImageFilter()
    f.SetReferenceImage(image)
    f.SetOutputOrigin(image.GetOrigin())
    f.SetOutputSpacing((spacing_x, spacing_y, spacing_z))
    f.SetSize((size_x, size_y, size_z))
    f.SetInterpolator(interpolator)
    f.SetOutputPixelType(sitk.sitkInt16)
    result = f.Execute(image)
    #result_around = sitk.GetImageFromArray(
    #    np.around(sitk.GetArrayFromImage(result)))
    return result


img_folder =  main_path + '/nifti/'
save_folder = save_path + '/nifti/'

label_folder =  main_path + 'labels/'
label_save_folder = save_path + 'labels/'
    
if not os.path.isdir(save_folder):
    os.makedirs(save_folder)
    os.makedirs(label_save_folder)
    
    
patients = [f[:-n] for f in os.listdir(img_folder) if f.endswith(end)]

for patient in tqdm(patients):

    ct = load_patient(patient, img_folder)
    spacing = ct.GetSpacing()
    shape = ct.GetSize()
    ct = do_resampling(image=ct, spacing_x=2, spacing_y=2, spacing_z=2, 
                       interpolator='Linear')

    # prefix name
    sitk.WriteImage(ct, save_folder + patient.lower() + end)
    
    try:
        mask = load_patient(patient, label_folder, is_label=True)
        mask = do_resampling(image=mask, spacing_x=2, spacing_y=2, 
                            spacing_z=2,
                            interpolator='NearestNeighbor')
    
        # prefix name
        sitk.WriteImage(mask, label_save_folder + patient.lower() + '-seg' + end)
    except:
        print('Patient {} has no labels'.format(patient))
