#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

WorkingDirectory = "/home/pkashyap/PK/PK2"

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

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

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

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

#Wraps command **flirt**
a_flirt = pe.MapNode(interface = fsl.FLIRT(), name = 'a_flirt', iterfield = ['in_file'])
a_flirt.inputs.reference = '/home/pkashyap/data/templates/MNI152_T1_1mm_first_brain_mask.nii'
a_flirt.inputs.cost = 'corratio'
a_flirt.inputs.interp = 'trilinear'
a_flirt.inputs.bins = 256
a_flirt.inputs.dof = 12
a_flirt.inputs.searchr_x = [-90,90]
a_flirt.inputs.searchr_y = [-90,90]
a_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.field_file = True
fnirt.inputs.jacobian_file = True
fnirt.inputs.warp_resolution = 10,10,10

#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'

#Wraps command **bbregister**
bbr = pe.MapNode(interface = freesurfer.BBRegister(), name = 'bbr', iterfield = ['subject_id', 'source_file'])
bbr.inputs.contrast_type = 't2'
bbr.inputs.init = 'fsl'
bbr.inputs.out_reg_file = 'bbregister.dof6.init.dat'
bbr.inputs.out_fsl_file = 'flirt.init.mtx'

#Wraps command **mri_vol2surf**
v2sl = pe.MapNode(interface = freesurfer.SampleToSurface(), name = 'v2sl', iterfield = ['source_file', 'reg_file', 'target_subject'])
v2sl.inputs.hemi = 'lh'
v2sl.inputs.sampling_range = .5
v2sl.inputs.sampling_units = 'frac'
v2sl.inputs.interp_method = 'trilinear'
v2sl.inputs.surf_reg = True
v2sl.inputs.args = '--cortex'

#Wraps command **mri_vol2surf**
v2sr = pe.MapNode(interface = freesurfer.SampleToSurface(), name = 'v2sr', iterfield = ['source_file', 'reg_file', 'target_subject'])
v2sr.inputs.hemi = 'rh'
v2sr.inputs.sampling_range = .5
v2sr.inputs.sampling_units = 'frac'
v2sr.inputs.interp_method = 'trilinear'
v2sr.inputs.surf_reg = True
v2sr.inputs.args = '--cortex'

#Wraps command **3dcalc**
calc3 = pe.MapNode(interface = afni.Calc(), name = 'calc3', iterfield = ['in_file_a'])
calc3.inputs.expr = 'step(a)'
#calc2.inputs.in_file_b = '/home/pkashyap/data/templates/MNI152_T1_1mm_first_brain_mask.nii'
calc3.inputs.outputtype = 'NIFTI'

#Wraps command **3dcalc**
calc4 = pe.MapNode(interface = afni.Calc(), name = 'calc4', iterfield = ['in_file_a'])
calc4.inputs.expr = 'step(a)'
#calc2.inputs.in_file_b = '/home/pkashyap/data/templates/MNI152_T1_1mm_first_brain_mask.nii'
calc4.inputs.outputtype = 'NIFTI'

#Wraps command **mri_surf2surf**
s2sl = pe.MapNode(interface = freesurfer.SurfaceSmooth(), name = 's2sl', iterfield = ['in_file', 'subject_id'])
s2sl.inputs.hemi = 'lh'
s2sl.inputs.fwhm = 2
s2sl.inputs.args = '--cortex'

#Wraps command **mri_surf2surf**
s2sr = pe.MapNode(interface = freesurfer.SurfaceSmooth(), name = 's2sr', iterfield = ['in_file', 'subject_id'])
s2sr.inputs.hemi = 'rh'
s2sr.inputs.fwhm = 2
s2sr.inputs.args = '--cortex'

#Wraps command **mri_vol2vol**
v2v = pe.MapNode(interface = freesurfer.ApplyVolTransform(), name = 'v2v', iterfield = ['source_file', 'reg_file'])
v2v.inputs.tal = True
v2v.inputs.tal_resolution = 2
v2v.inputs.interp = 'nearest'


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

#Wraps command **3dCalc**
tcat= pe.MapNode(interface = afni.Calc(), name = 'tcat', iterfield = ['in_file_a'])
tcat.inputs.start_idx = 5
tcat.inputs.stop_idx = 299
tcat.inputs.expr = 'a'
tcat.inputs.outputtype = 'NIFTI'

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

#Wraps command **3dDespike**
despike = pe.MapNode(interface = afni.Despike(), name = 'despike', iterfield = ['in_file'])
despike.inputs.args = '-NEW'
despike.inputs.outputtype = 'NIFTI'

#Wraps command **3dTshift**
tshift = pe.MapNode(interface = afni.TShift(), name = 'tshift', iterfield = ['in_file'])
tshift.inputs.tzero = 0
tshift.inputs.tpattern = 'seqplus'
tshift.inputs.outputtype = 'NIFTI'

#Wraps command **3dTstat**
tstat = pe.MapNode(interface = afni.TStat(), name = 'tstat', iterfield = ['in_file'])
tstat.inputs.options = '-mean'
tstat.inputs.outputtype = 'NIFTI'

#Wraps command **3dvolreg**
volreg = pe.MapNode(interface = afni.Volreg(), name = 'volreg', iterfield = ['in_file', 'basefile'])
volreg.inputs.zpad = 4
volreg.inputs.args = '-Fourier -twopass'
volreg.inputs.outputtype = 'NIFTI'


#Wraps command **3dcalc**
calc2 = pe.MapNode(interface = afni.Calc(), name = 'calc2', iterfield = ['in_file_a'])
calc2.inputs.expr = 'step(a)'
#calc2.inputs.in_file_b = '/home/pkashyap/data/templates/MNI152_T1_1mm_first_brain_mask.nii'
calc2.inputs.outputtype = 'NIFTI'

#Wraps command **flirt**
af_flirt = pe.MapNode(interface = fsl.FLIRT(), name = 'af_flirt', iterfield = ['in_file', 'reference'])
af_flirt.inputs.cost = 'corratio'
af_flirt.inputs.interp = 'trilinear'
af_flirt.inputs.dof = 6

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

#Create a workflow to connect all those nodes
analysisflow = nipype.Workflow('MyWorkflow')
analysisflow.connect(io_sf, 'anat', a_refit, 'in_file')
analysisflow.connect(a_refit, 'out_file', a_resam, 'in_file')
analysisflow.connect(a_resam, 'out_file', bet, 'in_file')
analysisflow.connect(bet, 'out_file', datasink, 'skullstrip')
analysisflow.connect(bet, 'out_file', a_flirt, 'in_file')
analysisflow.connect(bet, 'out_file', fnirt, 'in_file')
analysisflow.connect(a_flirt, 'out_matrix_file', fnirt, 'affine_file')
analysisflow.connect(a_flirt, 'out_matrix_file', datasink, 'aff_mat')
analysisflow.connect(fnirt, 'field_file', datasink, 'mat_f')
analysisflow.connect(fnirt, 'jacobian_file', datasink, 'mat_j')
analysisflow.connect(fnirt, 'fieldcoeff_file', datasink, 'mat_c')
analysisflow.connect(fnirt, 'warped_file', datasink, 'warp')
analysisflow.connect(io_sf, 'func', f_refit, 'in_file')
analysisflow.connect(f_refit, 'out_file', tcat, 'in_file_a')
analysisflow.connect(tcat, 'out_file', f_resam, 'in_file')
analysisflow.connect(f_resam, 'out_file', despike, 'in_file')
analysisflow.connect(despike, 'out_file', tshift, 'in_file')
analysisflow.connect(tshift, 'out_file', tstat, 'in_file')
analysisflow.connect(tshift, 'out_file', volreg, 'in_file')
analysisflow.connect(tstat, 'out_file', volreg, 'basefile')
analysisflow.connect(volreg, 'out_file', af_flirt, 'in_file')
analysisflow.connect(volreg, 'oned_file', datasink, 'vr_1d')
analysisflow.connect(volreg, 'out_file', datasink, 'func_vr')
analysisflow.connect(fnirt, 'warped_file', af_flirt, 'reference')
analysisflow.connect(af_flirt, 'out_file', datasink, 'func_flirt')
analysisflow.connect(af_flirt, 'out_matrix_file', datasink, 'func_aff_mat')
analysisflow.connect(a_resam, 'out_file', recon, 'T1_files')
analysisflow.connect(recon, 'subjects_dir', datasink, 'segu')
analysisflow.connect(recon, 'subject_id', datasink, 's_id')
analysisflow.connect(recon, 'subject_id', bbr, 'subject_id')
analysisflow.connect(tstat, 'out_file', bbr, 'source_file')
analysisflow.connect(bbr, 'out_fsl_file', datasink, 'bbr_fslmat')
analysisflow.connect(bbr, 'out_reg_file', datasink, 'bbr_reg')
analysisflow.connect(tstat, 'out_file', calc2, 'in_file_a')
analysisflow.connect(calc2, 'out_file', v2sl, 'source_file')
analysisflow.connect(bbr, 'out_reg_file', v2sl, 'reg_file')
analysisflow.connect(recon, 'subjects_dir', v2sl, 'target_subject')
analysisflow.connect(v2sl, 'out_file', datasink, 'v2sl')
analysisflow.connect(calc2, 'out_file', v2sr, 'source_file')
analysisflow.connect(bbr, 'out_reg_file', v2sr, 'reg_file')
analysisflow.connect(recon, 'subjects_dir', v2sr, 'target_subject')
analysisflow.connect(v2sr, 'out_file', datasink, 'v2sr')
analysisflow.connect(v2sl, 'out_file', s2sl, 'in_file')
analysisflow.connect(recon, 'subject_id', s2sl, 'subject_id')
analysisflow.connect(s2sl, 'out_file', calc3, 'in_file_a')
analysisflow.connect(v2sr, 'out_file', s2sr, 'in_file')
analysisflow.connect(recon, 'subject_id', s2sr, 'subject_id')
analysisflow.connect(s2sr, 'out_file', calc4, 'in_file_a')
analysisflow.connect(calc4, 'out_file', datasink, 's2s_m_r')
analysisflow.connect(calc3, 'out_file', datasink, 's2s_m_l')
analysisflow.connect(calc2, 'out_file', v2v, 'source_file')
analysisflow.connect(bbr, 'out_reg_file', v2v, 'reg_file')
analysisflow.connect(v2v, 'transformed_file', datasink, 'v2v_t')


#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)