#!/usr/bin/env python
"""
=============================================
dMRI: Connectivity - Camino, CMTK, FreeSurfer
=============================================

Introduction
============

This script, connectivity_tutorial.py, demonstrates the ability to perform connectivity mapping
using Nipype for pipelining, Freesurfer for Reconstruction / Parcellation, Camino for tensor-fitting
and tractography, and the Connectome Mapping Toolkit (CMTK) for connectivity analysis.

    python connectivity_tutorial.py

We perform this analysis using the FSL course data, which can be acquired from here:

    * http://www.fmrib.ox.ac.uk/fslcourse/fsl_course_data2.tar.gz

This pipeline also requires the Freesurfer directory for 'subj1' from the FSL course data.
To save time, this data can be downloaded from here:

    * http://dl.dropbox.com/u/315714/subj1.zip?dl=1

A data package containing the outputs of this pipeline can be obtained from here:

    * http://db.tt/1vx4vLeP

Along with Camino (http://web4.cs.ucl.ac.uk/research/medic/camino/pmwiki/pmwiki.php?n=Main.HomePage),
Camino-Trackvis (http://www.nitrc.org/projects/camino-trackvis/), FSL (http://www.fmrib.ox.ac.uk/fsl/),
and Freesurfer (http://surfer.nmr.mgh.harvard.edu/), you must also have the Connectome File Format
library installed as well as the Connectome Mapper.

These are written by Stephan Gerhard and can be obtained from:

    http://www.cmtk.org/

Or on github at:

    CFFlib: https://github.com/LTS5/cfflib
    CMP: https://github.com/LTS5/cmp

Output data can be visualized in the ConnectomeViewer

    ConnectomeViewer: https://github.com/LTS5/connectomeviewer

First, we import the necessary modules from nipype.
"""

import nipype.interfaces.io as nio           # Data i/o
import nipype.interfaces.utility as util     # utility
import nipype.pipeline.engine as pe          # pypeline engine
import nipype.interfaces.camino as camino
import nipype.interfaces.fsl as fsl
import nipype.interfaces.camino2trackvis as cam2trk
import nipype.interfaces.freesurfer as fs    # freesurfer
import nipype.interfaces.cmtk as cmtk
import nipype.algorithms.misc as misc
import inspect

import os.path as op                      # system functions
import cmp                                    # connectome mapper

"""
We define the following functions to scrape the voxel and data dimensions of the input images. This allows the
pipeline to be flexible enough to accept and process images of varying size. The SPM Face tutorial
(fmri_spm_face.py) also implements this inferral of voxel size from the data. We also define functions to
select the proper parcellation/segregation file from Freesurfer's output for each subject. For the mapping in
this tutorial, we use the aparc+seg.mgz file. While it is possible to change this to use the regions defined in
aparc.a2009s+aseg.mgz, one would also have to write/obtain a network resolution map defining the nodes based on those
regions.
"""

def get_vox_dims(volume):
    import nibabel as nb
    if isinstance(volume, list):
        volume = volume[0]
    nii = nb.load(volume)
    hdr = nii.get_header()
    voxdims = hdr.get_zooms()
    return [float(voxdims[0]), float(voxdims[1]), float(voxdims[2])]

def get_data_dims(volume):
    import nibabel as nb
    if isinstance(volume, list):
        volume = volume[0]
    nii = nb.load(volume)
    hdr = nii.get_header()
    datadims = hdr.get_data_shape()
    return [int(datadims[0]), int(datadims[1]), int(datadims[2])]

def get_affine(volume):
    import nibabel as nb
    nii = nb.load(volume)
    return nii.get_affine()

def select_aparc(list_of_files):
    for in_file in list_of_files:
        if 'aparc+aseg.mgz' in in_file:
            idx = list_of_files.index(in_file)
    return list_of_files[idx]

def select_aparc_annot(list_of_files):
    for in_file in list_of_files:
        if '.aparc.annot' in in_file:
            idx = list_of_files.index(in_file)
    return list_of_files[idx]

"""
These need to point to the main Freesurfer directory as well as the freesurfer subjects directory.
No assumptions are made about where the directory of subjects is placed.
Recon-all must have been run on subj1 from the FSL course data.
"""

fs_dir = op.abspath('/usr/local/freesurfer')
subjects_dir = op.abspath(op.join(op.curdir,'./subjects'))
fsl.FSLCommand.set_default_output_type('NIFTI')

"""
This needs to point to the fdt folder you can find after extracting
http://www.fmrib.ox.ac.uk/fslcourse/fsl_course_data2.tar.gz
"""

data_dir = op.abspath('fsl_course_data/fdt/')
fs.FSCommand.set_default_subjects_dir(subjects_dir)
subject_list = ['subj1']

"""
An infosource node is used to loop through the subject list and define the input files.
For our purposes, these are the diffusion-weighted MR image, b vectors, and b values.
The info dictionary is used to provide a template of the naming of these files. For instance,
the 4D nifti diffusion image is stored in the FSL course data as data.nii.gz.
"""

infosource = pe.Node(interface=util.IdentityInterface(fields=['subject_id']), name="infosource")
infosource.iterables = ('subject_id', subject_list)

info = dict(dwi=[['subject_id', 'data']],
            bvecs=[['subject_id','bvecs']],
            bvals=[['subject_id','bvals']])

"""
A datasource node is used to perform the actual data grabbing.
Templates for the associated images are used to obtain the correct images.
The data are assumed to lie in data_dir/subject_id/.
"""

datasource = pe.Node(interface=nio.DataGrabber(infields=['subject_id'],
                                               outfields=info.keys()),
                     name = 'datasource')

datasource.inputs.template = "%s/%s"
datasource.inputs.base_directory = data_dir
datasource.inputs.field_template = dict(dwi='%s/%s.nii.gz')
datasource.inputs.template_args = info
datasource.inputs.base_directory = data_dir
datasource.inputs.sort_filelist = True

"""
FreeSurferSource nodes are used to retrieve a number of image
files that were automatically generated by the recon-all process.
Here we use three of these nodes, two of which are defined to return files for solely the left and right hemispheres.
"""

FreeSurferSource = pe.Node(interface=nio.FreeSurferSource(), name='fssource')
FreeSurferSource.inputs.subjects_dir = subjects_dir

FreeSurferSourceLH = pe.Node(interface=nio.FreeSurferSource(), name='fssourceLH')
FreeSurferSourceLH.inputs.subjects_dir = subjects_dir
FreeSurferSourceLH.inputs.hemi = 'lh'

FreeSurferSourceRH = pe.Node(interface=nio.FreeSurferSource(), name='fssourceRH')
FreeSurferSourceRH.inputs.subjects_dir = subjects_dir
FreeSurferSourceRH.inputs.hemi = 'rh'

"""
Since the b values and b vectors come from the FSL course, we must convert it to a scheme file
for use in Camino.
"""

fsl2scheme = pe.Node(interface=camino.FSL2Scheme(), name="fsl2scheme")
fsl2scheme.inputs.usegradmod = True

"""
FSL's Brain Extraction tool is used to create a mask from the b0 image
"""

b0Strip = pe.Node(interface=fsl.BET(mask = True), name = 'bet_b0')

"""
FSL's FLIRT function is used to coregister the b0 mask and the structural image.
A convert_xfm node is then used to obtain the inverse of the transformation matrix.