%pylab inline

import numpy as np
import matplotlib.pyplot as plt

import os
import nibabel as nib

MNI_PATH = 'mni_icbm152_nlin_asym_09a'

gm_fname = os.path.join(MNI_PATH, 'mni_icbm152_gm_tal_nlin_asym_09a.nii')
wm_fname = os.path.join(MNI_PATH, 'mni_icbm152_wm_tal_nlin_asym_09a.nii')
csf_fname = os.path.join(MNI_PATH, 'mni_icbm152_csf_tal_nlin_asym_09a.nii')

gm_img = nib.load(gm_fname)
gm_data = gm_img.get_data()
wm_img = nib.load(wm_fname)
wm_data = wm_img.get_data()
csf_img = nib.load(csf_fname)
csf_data = csf_img.get_data()
csf_data.shape

t1_fname = os.path.join(MNI_PATH, 'mni_icbm152_t1_tal_nlin_asym_09a.nii')
t2_fname = os.path.join(MNI_PATH, 'mni_icbm152_t2_tal_nlin_asym_09a.nii')

t1_slice = nib.load(t1_fname).get_data()[:, :, 94]
t2_slice = nib.load(t2_fname).get_data()[:, :, 94]
plt.imshow(np.hstack((t1_slice, t2_slice)), cmap="gray")

# Get binary slices for gray matter, white matter, CSF
mid_gm = gm_data[:, :, 94] > 0.5
mid_wm = wm_data[:, :, 94] > 0.5
mid_csf = csf_data[:, :, 94] > 0.5

plt.imshow(np.hstack((mid_gm, mid_wm, mid_csf)), cmap="gray")

# Combine GM, WM, CSF in different proportions for T1-like and T2-like
t1_like = mid_gm * 0.7 + mid_wm * 1.0 + mid_csf * 0.1
t2_like = mid_gm * 0.5 + mid_wm * 0.2 + mid_csf * 0.8
plt.imshow(np.hstack((t1_like, t2_like)), cmap="gray")

plt.plot(t1_like.ravel(), t2_like.ravel(), '.')
plt.axis([-0.1, 1.1, -0.1, 1.1])

hgram_registered, x_edges, y_edges = np.histogram2d(t1_like.ravel(), t2_like.ravel(), 30)
hgram_registered.max()

plt.imshow(np.clip(hgram_registered, 0, 1000), cmap='gray', interpolation='nearest')

def mutual_information(hgram):
    """ Mutual information for joint histogram
    """
    pxy = hgram / float(np.sum(hgram)) # Convert to probability
    px = np.sum(pxy, 1) # marginal for x over y
    py = np.sum(pxy, 0) # marginal for y over x
    px_py = px[:, None] * py[None, :] # Broadcast to multiply marginals
    nzs = pxy > 0 # Only non-zero pxy values contribute to the sum
    return np.sum(pxy[nzs] * np.log(pxy[nzs] / px_py[nzs]))

mutual_information(hgram_registered)

t2_like_moved = np.zeros_like(t2_like)
t2_like_moved[10:, :] = t2_like[:-10, :]
plt.imshow(t2_like_moved, cmap="gray")

plt.plot(t1_like.ravel(), t2_like_moved.ravel(), '.')
plt.axis([-0.1, 1.1, -0.1, 1.1])

hgram_off, x_edges, y_edges = np.histogram2d(t1_like.ravel(), t2_like_moved.ravel(), 30)
plt.imshow(np.clip(hgram_off, 0, 1000), cmap='gray', interpolation='nearest')

mutual_information(hgram_off)