# - compatibility with Python 3
from __future__ import print_function  # print('me') instead of print 'me'
from __future__ import division  # 1/2 == 0.5, not 0

# - show figures inside the notebook
%matplotlib inline

# - import common modules
import numpy as np  # the Python array package
import matplotlib.pyplot as plt  # the Python plotting package

# - set gray colormap and nearest neighbor interpolation by default
plt.rcParams['image.cmap'] = 'gray'
plt.rcParams['image.interpolation'] = 'nearest'

import nibabel as nib
img = nib.load('ds107_sub012_t1r2.nii')
data = img.get_data()
data.shape

data = data.astype(np.float32)

vol0 = data[..., 0]
vol1 = data[..., 1]
vol1.shape

mid_vol0 = vol0[:, :, 17]
plt.imshow(mid_vol0)

mid_vol1 = vol1[:, :, 17]
plt.imshow(mid_vol1)

fig, axes = plt.subplots(1, 2, figsize=(10, 20))
axes[0].imshow(mid_vol0)
axes[0].set_title('slice 0')
axes[1].imshow(mid_vol1)
axes[1].set_title('slice 1')

# Make slice full of zeros like mid_vol1 slice
shifted_mid_vol1 = np.zeros(mid_vol1.shape)
# Fill the lower 54 (of 64) x lines with mid_vol1
shifted_mid_vol1[8:, :] = mid_vol1[:-8, :]
# Now we have something like mid_vol1 but translated in the first dimension
plt.imshow(shifted_mid_vol1)

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
    """
    # Make a 0-filled array of same shape as `img_slice`
    trans_slice = np.zeros(img_slice.shape)
    # Use slicing to select voxels out of the image and move them
    # Up or down on the first (x) axis
    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

plt.imshow(mid_vol0 - shifted_mid_vol1)

def mean_abs_mismatch(slice0, slice1):
    """ Mean absoute difference between images
    """
    return np.mean(np.abs(slice0 - slice1))

mismatches = []
translations = range(-25, 15)  # Candidate values for t
for t in translations:
    # Make the translated image Y_t
    unshifted = x_trans_slice(shifted_mid_vol1, t)
    # Calculate the mismatch
    mismatch = mean_abs_mismatch(unshifted, mid_vol0)
    # Store it for later
    mismatches.append(mismatch)

plt.plot(translations, mismatches)
plt.xlabel('translation (t)')
plt.ylabel('mean absolute difference')

# Number of voxels in the image
n_voxels = np.prod(mid_vol1.shape)
# Reshape vol0 slice as 1D vector
mid_vol0_as_1d = mid_vol0.reshape(n_voxels)
# Reshape vol1 slice as 1D vector
mid_vol1_as_1d = mid_vol1.reshape(n_voxels)
# These original slices should be very close to each other already
# So - plotting one set of image values against the other should
# be close to a straight line
plt.plot(mid_vol0_as_1d, mid_vol1_as_1d, '.')
plt.xlabel('voxels in vol0 slice')
plt.ylabel('voxels in original vol1 slice')
# Correlation coefficient between them
np.corrcoef(mid_vol0_as_1d, mid_vol1_as_1d)[0, 1]

# The shifted slice will be less well matched
# Therefore the line will be less straight and narrow
plt.plot(mid_vol0_as_1d, shifted_mid_vol1.ravel(), '.')
plt.xlabel('voxels in vol0 slice')
plt.ylabel('voxels in shifted vol1 slice')
# Correlation coefficient between them will be nearer 0
np.corrcoef(mid_vol0_as_1d, shifted_mid_vol1.ravel())[0, 1]

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

correl_mismatches = []
translations = range(-25, 15)  # Candidate values for t
for t in translations:
    unshifted = x_trans_slice(shifted_mid_vol1, t)
    mismatch = correl_mismatch(unshifted, mid_vol0)
    correl_mismatches.append(mismatch)

plt.plot(translations, correl_mismatches)
plt.xlabel('translation (t)')
plt.ylabel('correlation mismatch')

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.
    """
    # Resample image using bilinear interpolation (order=1)
    trans_slice = snd.affine_transform(img_slice, [1, 1], [-x_vox_trans, 0], order=1)
    return trans_slice

fine_mismatches = []
fine_translations = np.linspace(-25, 15, 100)
for t in fine_translations:
    unshifted = fancy_x_trans_slice(shifted_mid_vol1, t)
    mismatch = correl_mismatch(unshifted, mid_vol0)
    fine_mismatches.append(mismatch)

plt.plot(fine_translations, fine_mismatches)

def cost_function(x_trans):
    # Function can use image slices defined in the global (notebook) scope
    # Calculate X_t - image translated by x_trans
    unshifted = fancy_x_trans_slice(shifted_mid_vol1, x_trans)
    # Return mismatch measure for the translated image X_t
    return correl_mismatch(unshifted, mid_vol0)

# value of the negative correlation for no translatino
cost_function(0)

# value of the negative correlation for translation of -8 voxels
cost_function(-8)

from scipy.optimize import fmin_powell

fmin_powell(cost_function, [0])

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

best_params = fmin_powell(cost_function, [0], callback=my_callback)
best_params

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)
    # Resample image using bilinear interpolation (order=1)
    trans_slice = snd.affine_transform(img_slice, [1, 1], -x_y_trans, order=1)
    return trans_slice

def fancy_cost_at_xy(x_y_trans):
    """ Give cost function at xy translation values `x_y_trans`
    """
    unshifted = fancy_xy_trans_slice(shifted_mid_vol1, x_y_trans)
    return correl_mismatch(unshifted, mid_vol0)

best_params = fmin_powell(fancy_cost_at_xy, [0, 0], callback=my_callback)
best_params

shifted_vol1 = np.zeros(vol1.shape)
shifted_vol1[8:, 5:, :] = vol1[:-8, :-5, :]
plt.imshow(shifted_vol1[:, :, 17])

def xyz_trans_vol(vol, x_y_z_trans):
    """ Make a new copy of `vol` translated by `x_y_z_trans` voxels
    
    x_y_z_trans is a sequence or array length 3, containing the (x, y, z) translations in voxels.
    
    Values in `x_y_z_trans` can be positive or negative, and can be floats.
    """
    x_y_z_trans = np.array(x_y_z_trans)
    # [1, 1, 1] says to do no zooming or rotation
    # Resample image using trilinear interpolation (order=1)
    trans_vol = snd.affine_transform(vol, [1, 1, 1], -x_y_z_trans, order=1)
    return trans_vol

def cost_at_xyz(x_y_z_trans):
    """ Give cost function value at xyz translation values `x_y_z_trans`
    """
    unshifted = xyz_trans_vol(shifted_vol1, x_y_z_trans)
    return correl_mismatch(unshifted, vol0)

best_params = fmin_powell(cost_at_xyz, [0, 0, 0], callback=my_callback)
best_params

unshifted_vol1 = snd.affine_transform(shifted_vol1, [1, 1, 1], -best_params)

fig, axes = plt.subplots(1, 2, figsize=(10, 20))
axes[0].imshow(vol0[:, :, 17])
axes[0].set_title('vol0')
axes[1].imshow(unshifted_vol1[:, :, 17])
axes[1].set_title('unshifted vol1')

correl_mismatches = []
translations = range(-25, 15)  # Candidate values for t
for t in translations:
    unshifted = x_trans_slice(shifted_mid_vol1, t)
    mismatch = correl_mismatch(unshifted, mid_vol0)
    correl_mismatches.append(mismatch)

plt.plot(translations, correl_mismatches)
plt.title('Cost as a function of $t$')

correl_mismatches = []
translations = range(-60, 50)  # Candidate values for t
for t in translations:
    unshifted = x_trans_slice(shifted_mid_vol1, t)
    mismatch = correl_mismatch(unshifted, mid_vol0)
    correl_mismatches.append(mismatch)

plt.plot(translations, correl_mismatches)
plt.title('Cost as a function of more $t$')

def cost_function(x_trans):
    # Function can use image slices defined in the global (notebook) scope
    # Calculate X_t - image translated by x_trans
    unshifted = fancy_x_trans_slice(shifted_mid_vol1, x_trans)
    # Return mismatch measure for the translated image X_t
    return correl_mismatch(unshifted, mid_vol0)

fmin_powell(cost_function, [35])