%pylab inline
import numpy as np
import matplotlib.pylab as plt
import os
import nibabel as nib

# assume we have the current directory the "mni_icbm152_nlin_asym_09a" directory that contains the MNI template
# if not, this can be downloaded from : http://www.bic.mni.mcgill.ca/ServicesAtlases/ICBM152NLin2009 - 
# choose the ICBM 2009a Nonlinear Asymmetric 1×1x1mm template
MNI_PATH = 'mni_icbm152_nlin_asym_09a'
os.listdir(MNI_PATH)

mni_fname = os.path.join(MNI_PATH, 'mni_icbm152_t1_tal_nlin_asym_09a.nii')
mni_img = nib.load(mni_fname)
plt.imshow(mni_img.get_data()[:, :, 30], cmap="gray")
print mni_img.get_affine()
print mni_img.shape

# where is our image to register:
DATA_PATH = os.path.join(os.path.expanduser('~'), 'data', 'ds105', 'sub001', 'anatomy')
os.listdir(DATA_PATH)

anat_fname = os.path.join(DATA_PATH, 'highres001.nii.gz')
print anat_fname

anat_img = nib.load(anat_fname)
plt.imshow(anat_img.get_data()[:, 70, :], cmap="gray")
print anat_img.get_affine()
print anat_img.shape

anat_hdr = anat_img.get_header()
print(anat_hdr)

#- Make sure we have pna-utils on the path
!echo $PYTHONPATH
#- import the fixing tool
import fixnifti

fixnifti.fixup_nifti_file(anat_fname)
fixed_fname = os.path.join(DATA_PATH, 'fhighres001.nii.gz')

# check that the new header contains information for the srow_{x,y,z}
fixed_anat = nib.load(fixed_fname)
fixed_anat_hdr = fixed_anat.get_header()
print(fixed_anat_hdr)

#save the fixed image in the current directory, call it 'my_fixed_anat.nii'  
nib.save(fixed_anat, 'my_fixed_anat.nii')


# make a nibabel image 
PNA_PATH = '/home/jb/code/pna-notebooks'
fixed_fname = os.path.join(PNA_PATH, 'my_fixed_anat.nii')
fixed_img = nib.load(fixed_fname)

!ants.sh

os.environ['ANTSPATH'] = '/home/jb/bin/ants/'

# RUN THIS IF YOU HAVE TIME ... 
# Output images will be prefixed with nrmlzd_"

# !ants.sh 3 my_fixed_anat.nii mni_icbm152_nlin_asym_09a/mni_icbm152_t1_tal_nlin_asym_09a.nii nrmlzd_

# we use viz from nipy:
from nipy.labs import viz

normed_img_fname = os.path.join(PNA_PATH, 'nrmlzd_deformed.nii.gz')
normed_img = nib.load(normed_img_fname)

slicer = viz.plot_anat(normed_img.get_data(), normed_img.get_affine(), dim=.2, black_bg=True)
slicer.edge_map(mni_img.get_data(), mni_img.get_affine())

def show_imgs(imgs, pos, titles=None, col="gray", interp="nearest"):
    """
    """
    
    Nimg = len(imgs)
    
    if (titles is None):
        titles = ['No title']*Nimg 
    
    if not isinstance(pos, list):
        pos = [pos]*Nimg
        
    fig, ax = plt.subplots(1,Nimg)
    
    for ix,im in enumerate(imgs):
        ax[ix].imshow(im.get_data()[pos[ix]], cmap=col)
        ax[ix].get_xaxis().set_ticks([])
        ax[ix].get_yaxis().set_ticks([])
        ax[ix].set_title(titles[ix])

pos = (slice(None), slice(None), 64)
show_imgs([mni_img, normed_img], [pos,pos], ['mni_img', 'normed_img'])
print normed_img.get_affine(), normed_img.shape
print mni_img.get_affine(), mni_img.shape

#PNA_PATH = '/home/jb/code/pna-notebooks'
#fixed_fname = os.path.join(PNA_PATH, 'my_fixed_anat.nii')
#fixed_img = nib.load(fixed_fname)

data = fixed_img.get_data()
aff = fixed_img.get_affine().copy()

# cut from z=40 to z=225
lz, hz = 40, 225

data_cut_neck = data[:,:,lz:hz]
aff[2,3] = -(hz-lz)/2. # put the origin in the middle of the z
print aff

cut_neck_fname = os.path.join(PNA_PATH,"neck_cut.nii.gz")
cut_neck_img = nib.Nifti1Pair(data_cut_neck, aff)

nib.save(cut_neck_img, cut_neck_fname)

# re-read and display
new_img = nib.load(cut_neck_fname)
plt.imshow(new_img.get_data()[:, 80, :], cmap="gray")
print new_img.get_affine()
print new_img.shape

# get help with : !antsIntroduction.sh -h
# launch it with option : -t RI (for Rigid) and prefix the results with "rigid_"

!antsIntroduction.sh -d 3 -i neck_cut.nii.gz -r mni_icbm152_nlin_asym_09a/mni_icbm152_t1_tal_nlin_asym_09a.nii  -o rigid_ -t RI

# !ANTS 3 -i fixed_anat.nii -r mni_icbm152_nlin_asym_09a/mni_icbm152_t1_tal_nlin_asym_09a.nii -o rigid_ -t RI
#!ANTS 3 -i neck_cut.nii.gz -r mni_icbm152_nlin_asym_09a/mni_icbm152_t1_tal_nlin_asym_09a.nii -o nrmlzd_ -s MI -t RI 
#--number-of-affine-iterations 10000x10000x10000

# look at the results:
rigid_fname = os.path.join(PNA_PATH, 'rigid_deformed.nii.gz')
rigid_img = nib.load(rigid_fname)

slicer = viz.plot_anat(rigid_img.get_data(), rigid_img.get_affine(), dim=.2, black_bg=True)
slicer.edge_map(mni_img.get_data(), mni_img.get_affine())

pos = (slice(None), slice(None), 64)
show_imgs([mni_img, rigid_img], [pos,pos], ['mni_img', 'rigid_img'])

# re-try with the rigid transform as the input
# Use antsIntroduction that gives us more flexibility : the -m 10x30x5 states the number of iterations
# by default 30x90x20

# !antsIntroduction.sh -d 3 -i rigid_deformed.nii.gz -r mni_icbm152_nlin_asym_09a/mni_icbm152_t1_tal_nlin_asym_09a.nii  \
#                    -o from_rigid_ -t SY -m 10x30x5


renormed_fname = os.path.join(PNA_PATH, 'from_rigid_deformed.nii.gz')
renormed_img = nib.load(renormed_fname)

slicer = viz.plot_anat(renormed_img.get_data(), renormed_img.get_affine(), dim=.2, black_bg=True)
slicer.edge_map(mni_img.get_data(), mni_img.get_affine())

pos = (slice(None),40,slice(None))
show_imgs([mni_img, renormed_img], pos, ['mni_img', 'renormed_img']) 

# load the deformation maps
# use slice objects

warp_fname = os.path.join(PNA_PATH, 'from_rigid_Warp.nii.gz')
warp_img = nib.load(warp_fname)
print warp_img.shape
print renormed_img.shape
print mni_img.shape

# The warp image is 5 dimensional : TO DO : EXPLAIN ... 
pos = (slice(None),50,slice(None))
pos = [pos, pos + (0, 0), pos]
show_imgs([renormed_img, warp_img, mni_img], pos, ['renormed', 'warp', 'mni'])


# increase regularisation using the following command:
#
# Compare the results with the previous normalized file

renormed_25_fname = os.path.join(PNA_PATH, 'from_rigid_25_deformed.nii.gz')
renormed_25_img = nib.load(renormed_25_fname)

pos = (slice(None),50,slice(None))
show_imgs([renormed_img, renormed_25_img, rigid_img], [pos, pos, pos], ['full elastic', '.25 regul', 'rigid']) 

renormed_01_fname = os.path.join(PNA_PATH, 'from_rigid_01_deformed.nii.gz')
renormed_01_img = nib.load(renormed_01_fname)

pos = (slice(None),50,slice(None))
show_imgs([mni_img, renormed_img, renormed_01_img, rigid_img], 
          [pos, pos, pos, pos], 
          ['MNI', '0 Regul', '.01 Regul', 'Rigid'])