%pylab inline

import numpy as np
import matplotlib.pyplot as plt

import os
MNI_PATH = 'mni_icbm152_nlin_asym_09a'
os.listdir(MNI_PATH)

import nibabel as nib
t2_fname = os.path.join(MNI_PATH, 'mni_icbm152_t2_tal_nlin_asym_09a.nii')
t2_img = nib.load(t2_fname)
plt.imshow(t2_img.get_data()[:, :, 90], cmap="gray")

DATA_PATH = os.path.join(os.path.expanduser('~'), 'data', 'ds105')
os.listdir(DATA_PATH)                             

anat_fname = os.path.join(DATA_PATH, 'sub001', 'anatomy', 'highres001.nii.gz')
anat_img = nib.load(anat_fname)
plt.imshow(anat_img.get_data()[:, :, 128], cmap="gray")

import nipy.algorithms.registration as nar

# Make a transform that does nothing to use the image affines only to register the datasets
empty_transform = nar.Affine()
empty_transform.param

anat_resampled = nar.resample(anat_img, empty_transform, t2_img)

fig, axes = plt.subplots(1, 2)
fig.set_size_inches(10.5,18.5)
axes[0].imshow(anat_resampled.get_data()[:, :, 90], cmap="gray")
axes[1].imshow(t2_img.get_data()[:, :, 90], cmap="gray")

register_obj = nar.HistogramRegistration(from_img = t2_img, to_img = anat_img, similarity='cc')

transform = register_obj.optimize('affine')

anat_better_resampled = nar.resample(anat_img, transform, t2_img)

fig, axes = plt.subplots(1, 2)
fig.set_size_inches(10.5,18.5)
axes[0].imshow(anat_better_resampled.get_data()[:, :, 90], cmap="gray")
axes[1].imshow(t2_img.get_data()[:, :, 90], cmap="gray")