%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt

subj_dir = 'sub009'



import nibabel as nib
fname = 'ds107_sub001_run01.nii.gz'
img = nib.load(fname)
data = img.get_data()
data.shape

data = data.astype(np.float32)

mid_vol0 = data[:, :, 17, 0]
plt.imshow(mid_vol0, cmap='gray')

mid_vol1 = data[:, :, 17, 1]
plt.imshow(mid_vol1, cmap='gray')

moved_mid_vol1 = np.zeros_like(mid_vol1)
moved_mid_vol1.shape

moved_mid_vol1[10:, :] = mid_vol1[:-10, :]
plt.imshow(moved_mid_vol1, cmap='gray')

def x_trans_slice(img_slice, x_vox_trans):
    """ Make a new copy of `img_slice` translated by `x_vox_trans` voxels
    
    `x_vox_trans` can be positive or negative
    """
    trans_slice = np.zeros_like(img_slice)
    if x_vox_trans < 0:
        trans_slice[:x_vox_trans, :] = img_slice[-x_vox_trans:, :]
    elif x_vox_trans == 0:
        trans_slice[:, :] = img_slice
    else:
        trans_slice[x_vox_trans:, :] = img_slice[:-x_vox_trans, :]
    return trans_slice

# Test positive translation does what we did before
assert np.allclose(x_trans_slice(mid_vol1, 10), moved_mid_vol1)
# Check negative translation reverses what we did before
back_again = x_trans_slice(moved_mid_vol1, -10)
assert np.allclose(back_again[:-10, :], mid_vol1[:-10, :])
# Check zero translation works
assert np.allclose(x_trans_slice(mid_vol1, 0), mid_vol1)

plt.imshow(mid_vol0 - moved_mid_vol1, cmap='gray')

def my_cost_func(slice0, slice1):
    """ Calculate a measure of similarity between two image slices
    """
    return np.mean(np.abs(slice0 - slice1))

costs = []
translations = range(-25, 15)
for t in translations:
    unmoved = x_trans_slice(moved_mid_vol1, t)
    cost = my_cost_func(unmoved, mid_vol0)
    costs.append(cost)

plt.plot(translations, costs)

def my_add(x, y):
    return x + y

def my_mult(x, y):
    return x * y

print(my_add(2, 3))
print(my_mult(2, 3))

def apply_func(x, y, func):
    return func(x, y)

print(apply_func(2, 3, my_add))
print(apply_func(2, 3, my_mult))

def cost_at_t(t, slice0, slice1, cost_func):
    """ Give cost function at translation value `t` and function `cost_func`
    """
    unmoved = x_trans_slice(slice0, t)
    return cost_func(unmoved, slice1)

new_costs = []
for t in translations:
    cost = cost_at_t(t, moved_mid_vol1, mid_vol0, my_cost_func)
    new_costs.append(cost)
assert new_costs == costs     

plt.plot(mid_vol1.ravel(), mid_vol0.ravel(), '.')

def correl_cost(slice0, slice1):
    """ Negative correlation between the two images, flattened to 1D """
    correl = np.corrcoef(slice0.ravel(), slice1.ravel())[0, 1]
    return -correl

corr_costs = [cost_at_t(t, moved_mid_vol1, mid_vol0, correl_cost)
              for t in translations]

plt.plot(translations, corr_costs)

import scipy.ndimage as snd

def fancy_x_trans_slice(img_slice, x_vox_trans):
    """ Make a new copy of `img_slice` translated by `x_vox_trans` voxels
    
    `x_vox_trans` can be positive or negative, and can be a float.
    """
    trans_slice = snd.affine_transform(img_slice, [1, 1], [-x_vox_trans, 0])
    return trans_slice

for t in translations:
    assert np.allclose(fancy_x_trans_slice(moved_mid_vol1, t),
                       x_trans_slice(moved_mid_vol1, t))

def fancy_cost_at_t(t, slice0, slice1, cost_func):
    """ Give cost function at translation value `t` and function `cost_func`
    """
    unmoved = fancy_x_trans_slice(slice0, t)
    return cost_func(unmoved, slice1)

fine_costs = []
fine_translations = np.linspace(-25, 15, 100)
for t in fine_translations:
    cost = fancy_cost_at_t(t, moved_mid_vol1, mid_vol0, my_cost_func)
    fine_costs.append(cost)

plt.plot(fine_translations, fine_costs)

import scipy.optimize as sopt

sopt.fmin(fancy_cost_at_t, [0], args=(moved_mid_vol1, mid_vol0, my_cost_func))

def my_callback(params):
   print("Trying parameters " + str(params))

sopt.fmin(fancy_cost_at_t, [0], args=(moved_mid_vol1, mid_vol0, my_cost_func), callback=my_callback)

def fancy_xy_trans_slice(img_slice, x_y_trans):
    """ Make a new copy of `img_slice` translated by `x_y_trans` voxels
    
    x_y_trans is a sequence or array length 2, containing the (x, y) translations in voxels.
    
    Values in `x_y_trans` can be positive or negative, and can be floats.
    """
    x_y_trans = np.array(x_y_trans)
    trans_slice = snd.affine_transform(img_slice, [1, 1], -x_y_trans)
    return trans_slice

for t in translations:
    assert np.allclose(fancy_xy_trans_slice(moved_mid_vol1, (t, 0)),
                       x_trans_slice(moved_mid_vol1, t))
    assert np.allclose(fancy_xy_trans_slice(moved_mid_vol1, (0, t)),
                       x_trans_slice(moved_mid_vol1.T, t).T)

def fancy_cost_at_xy_t(xy_t, slice0, slice1, cost_func):
    """ Give cost function at xy translation values `xy_t` and function `cost_func`
    """
    unmoved = fancy_xy_trans_slice(slice0, xy_t)
    return cost_func(unmoved, slice1)

sopt.fmin(fancy_cost_at_xy_t, [0, 0], args=(moved_mid_vol1, mid_vol0, correl_cost), callback=my_callback)

anat_fname = 'ds107_sub001_highres.nii.gz'
anat_img = nib.load(anat_fname)
anat_data = anat_img.get_data()
anat_data.shape

data.shape

fig, axes = plt.subplots(1, 2)
fig.set_size_inches(10.5,18.5)
axes[0].imshow(anat_data[:, :, 96], cmap="gray")
axes[1].imshow(data[:, :, 17, 0], cmap="gray")

anat_img.get_affine()

img.get_affine()

funcvox2mm = img.get_affine()
anatvox2mm = anat_img.get_affine()

mm2anatvox = np.linalg.inv(anatvox2mm)

funcvox2anatvox = np.dot(mm2anatvox, funcvox2mm)
funcvox2anatvox

RZS = funcvox2anatvox[:3, :3]
translations = funcvox2anatvox[:3, 3]

resamp_anat = snd.affine_transform(anat_data, RZS, translations, output_shape=data.shape[:-1])
resamp_anat.shape

resamp_mid = resamp_anat[:, :, 17]
fig, axes = plt.subplots(1, 2)
fig.set_size_inches(10.5,18.5)
axes[0].imshow(resamp_mid, cmap="gray")
axes[1].imshow(data[:, :, 17, 0], cmap="gray")

plt.plot(resamp_mid.ravel(), mid_vol0.ravel(), '.')