from glob import glob
import os
from nipype import Workflow, MapNode, Node
from nipype.interfaces.ants import ApplyTransforms, AntsJointFusion, LabelGeometry
from nipype.utils.misc import human_order_sorted

T = glob('/data/out/ants_subjects/arno/antsTemplateToSubject*')[::-1]
ref = '/data/T1.nii.gz'
mask = '/data/out/ants_subjects/arno/antsBrainExtractionMask.nii.gz'

T1s = human_order_sorted(glob('/opt/data/OASIS-TRT_brains_to_OASIS_Atropos_template/*.nii.gz'))
labels = human_order_sorted(glob('/opt/data/OASIS-TRT_labels_to_OASIS_Atropos_template/*.nii.gz'))
thickness = '/data/out/ants_subjects/arno/antsCorticalThickness.nii.gz'
N = 20

wf = Workflow('labelflow')
transformer = MapNode(ApplyTransforms(), iterfield=['input_image'], name="transformer")
transformer.inputs.reference_image = ref
transformer.inputs.transforms = T
transformer.inputs.input_image = T1s[:N]
transformer.inputs.dimension = 3
transformer.inputs.invert_transform_flags = [False, False]

transformer_nn = MapNode(ApplyTransforms(), iterfield=['input_image'], name="transformer_nn")
transformer_nn.inputs.reference_image = ref
transformer_nn.inputs.transforms = T
transformer_nn.inputs.dimension = 3
transformer_nn.inputs.invert_transform_flags = [False, False]
transformer_nn.inputs.input_image = labels[:N]
transformer_nn.inputs.interpolation = 'NearestNeighbor'

labeler = Node(AntsJointFusion(), name='labeler')
labeler.inputs.dimension = 3
labeler.inputs.target_image = [ref]
labeler.inputs.out_label_fusion = 'label.nii.gz'
labeler.inputs.mask_image = mask
labeler.inputs.num_threads = 8
wf.connect(transformer, 'output_image', labeler, 'atlas_image')
wf.connect(transformer_nn, 'output_image', labeler, 'atlas_segmentation_image')

tocsv = Node(LabelGeometry(), name='get_measures')
tocsv.inputs.intensity_image = thickness
wf.connect(labeler, 'out_label_fusion', tocsv, 'label_image')

wf.base_dir = os.getcwd()  

wf.config['monitoring'] = {'enabled': True}
wf.run('MultiProc')