#This is a Nipype generator. Warning, here be dragons.
#!/usr/bin/env python
import sys
import nipype
import nipype.pipeline as pe
import nipype.interfaces.io as io
import nipype.interfaces.afni as afni
import nipype.interfaces.fsl as fsl
import nipype.interfaces.freesurfer as freesurfer
#from nipype import IdentityInterface

WorkingDirectory = "/home/pkashyap/PK/PK3"

#Flexibly collect data from disk to feed into workflows.
iop = pe.Node(io.SelectFiles(templates={'anat':'/home/pkashyap/data/subj{subj_id}/anat/t1{sess_id}_sanlm.nii.gz'}), name = 'iop')
iop.inputs.base_directory = '/home/pkashyap/PK/PK3'
iop.inputs.sort_filelist = True
iop.inputs.force_lists = True
iop.iterables = [('sess_id', ['1','2']), ('subj_id', ['01'])]


#Wraps command **3drefit**
refit = pe.MapNode(interface = afni.Refit(), name = 'refit', iterfield = ['in_file'])
refit.inputs.deoblique = True

#Wraps command **3dresample**
resam = pe.MapNode(interface = afni.Resample(), name = 'resam', iterfield = ['in_file'])
#resam2.inputs.orientation = 'RPI'
resam.inputs.master = '/home/pkashyap/data/templates/MNI152_T1_1mm_first_brain_mask.nii'
resam.inputs.outputtype = 'NIFTI'  

#Wraps command **bet**
bet = pe.MapNode(interface = fsl.BET(), name = 'bet', iterfield = ['in_file'])

#Wraps command **flirt**
flirt = pe.MapNode(interface = fsl.FLIRT(), name = 'flirt', iterfield = ['in_file'])
flirt.inputs.reference = '/home/pkashyap/data/templates/MNI152_T1_1mm_first_brain_mask.nii'
flirt.inputs.out_matrix_file = 'aff_{subj_id}_{sess_id}.mat'
flirt.inputs.cost = 'corratio'
flirt.inputs.interp = 'trilinear'
flirt.inputs.bins = 256
flirt.inputs.dof = 12
flirt.inputs.searchr_x = [-90,90]
flirt.inputs.searchr_y = [-90,90]
flirt.inputs.searchr_z = [-90,90]

#Wraps command **fnirt**
fnirt = pe.MapNode(interface = fsl.FNIRT(), name = 'fnirt', iterfield = ['in_file', 'affine_file'])
fnirt.inputs.ref_file = '/home/pkashyap/data/templates/MNI152_T1_1mm_first_brain_mask.nii'
fnirt.inputs.fieldcoeff_file = True
#fnirt.inputs.warped_file = 'highres.nii'
fnirt.inputs.field_file = True
fnirt.inputs.jacobian_file = True
fnirt.inputs.warp_resolution = 10,10,10


#identinterface
#infosource = pe.Node(IdentityInterface(fields=['subj_id','sess_id']), name='infosource')
#infosource.iterables = [('sess_id', ['1','2']), ('subj_id', ['01'])]

  

#Wraps command **recon-all**
recon = pe.MapNode(interface = freesurfer.ReconAll(), name = 'recon', iterfield = ['T1_files'])
recon.inputs.directive = 'all'
recon.inputs.openmp = 64
recon.inputs.subjects_dir = '.'
recon.inputs.subject_id = 'segu'
#recon.iterables = [('sess_id', ['1','2']), ('subj_id', ['01'])]



#Generic datasink module to store structured outputs
ds = pe.Node(interface = io.DataSink(), name = 'ds')
ds.inputs.base_directory = '/home/pkashyap/PK/PK3/results'

#Create a workflow to connect all those nodes
analysisflow = nipype.Workflow('MyWorkflow')

analysisflow.connect(iop, 'anat', refit, 'in_file')
analysisflow.connect(refit, 'out_file', resam, 'in_file')
analysisflow.connect(resam, 'out_file', bet, 'in_file')
analysisflow.connect(bet, 'out_file', flirt, 'in_file')
analysisflow.connect(flirt, 'out_matrix_file', fnirt, 'affine_file')
analysisflow.connect(bet, 'out_file', fnirt, 'in_file')
analysisflow.connect(iop, 'anat', recon, 'T1_files')
#analysisflow.connect(infosource,[('subj_id','sess_id')], recon,[('subj_id','sess_id')])
analysisflow.connect(fnirt, 'jacobian_file', ds, 'warp_mat_j')
analysisflow.connect(fnirt, 'field_file', ds, 'warp_mat_f')
analysisflow.connect(fnirt, 'fieldcoeff_file', ds, 'warp_mat_c')
analysisflow.connect(fnirt, 'warped_file', ds, 'warp')
analysisflow.connect(bet, 'out_file', ds, 'skullstrip')
analysisflow.connect(flirt, 'out_matrix_file', ds, 'aff_mat')
analysisflow.connect(recon, 'subjects_dir', ds, 'segu')
analysisflow.connect(recon, 'subject_id', ds, 's_id')


#Run the workflow
plugin = 'MultiProc' #adjust your desired plugin here
plugin_args = {'n_procs': 64} #adjust to your number of cores
#analysisflow.write_graph(graph2use='flat', format='png', simple_form=False)
analysisflow.run(plugin=plugin, plugin_args=plugin_args)