#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jun  3 10:46:04 2019

@author: Josefine
"""

import re
import glob
import sitk_functions as func
import SimpleITK as sitk

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)

filename_img = natural_sort(glob.glob('WHS/ct_train_test/ct_train/*_image.nii.gz')) # list of file names
filename_label = natural_sort(glob.glob('WHS/ct_train_test/ct_train/*_label.nii.gz')) # list of file names

for k in range(len(filename_img)):
    print('Load image',(k+1))
    img_sitk = sitk.ReadImage(filename_img[k])
    label_sitk = sitk.ReadImage(filename_label[k])
    filter = sitk.MinimumMaximumImageFilter()
    filter.Execute(img_sitk)
    min_val = filter.GetMinimum()
    
    # NORMAL
    sitk.WriteImage(img_sitk,'WHS/Augment_data/normal_{}_image.nii'.format(k))
    sitk.WriteImage(label_sitk,'WHS/Augment_data/normal_{}_label.nii'.format(k))

#    # ROTATION
#    [img_rot, label_rot] = func.affine_rotate(img_sitk,label_sitk,min_val)
#    sitk.WriteImage(img_rot,'WHS/Augment_data/rot_{}_image.nii'.format(k))
#    sitk.WriteImage(label_rot,'WHS/Augment_data/rot_{}_label.nii'.format(k))

#    # SHEAR
#    [img_sh, label_sh] = func.affine_shear(img_sitk,label_sitk,min_val)
#    sitk.WriteImage(img_sh,'WHS/Augment_data/sh_{}_image.nii'.format(k))
#    sitk.WriteImage(label_sh,'WHS/Augment_data/sh_{}_label.nii'.format(k))
#    
#    # Intensity
#    l1 = func.mult_and_add_intensity_fields(img_sitk)
#    sitk.WriteImage(l1,'WHS/Augment_data/intensity_{}_image.nii'.format(k))
#    sitk.WriteImage(label_sitk,'WHS/Augment_data/intensity_{}_label.nii'.format(k))

#    # B SPLINE
    numcontrolpoints = 5
    stdDeform = 15
    dim = 3
    [img_bspline, label_bspline] = func.BSplineDeform(img_sitk,label_sitk, dim, numcontrolpoints, stdDeform,min_val)
    sitk.WriteImage(img_bspline,'WHS/Augment_data/bspline_{}_image.nii'.format(k))
    sitk.WriteImage(label_bspline,'WHS/Augment_data/bspline_{}_label.nii'.format(k))

    # B SPLINE
    numcontrolpoints = 10
    stdDeform = 10
    dim = 3
    [img_bspline, label_bspline] = func.BSplineDeform(img_sitk,label_sitk, dim, numcontrolpoints, stdDeform,min_val)
    sitk.WriteImage(img_bspline,'WHS/Augment_data/bspline2_{}_image.nii'.format(k))
    sitk.WriteImage(label_bspline,'WHS/Augment_data/bspline2_{}_label.nii'.format(k))
    print('Finished with image',(k+1))