%pylab inline

import numpy as np
import matplotlib.pylab as plt

import nibabel as nib

fname = 'bold.nii.gz'

img = nib.load(fname)
img.shape

data = img.get_data()

data.dtype

data = data.astype(np.float32) # We need higher precision for our calculations

difference_volumes = np.diff(data, axis=3)
difference_volumes.shape

sq_diff = difference_volumes ** 2

n_voxels = np.prod(sq_diff.shape[:-1])
n_scans = sq_diff.shape[-1]
n_voxels, n_scans

vox_by_scans = np.reshape(sq_diff, (n_voxels, n_scans))

vox_by_scans.shape

scan_means = np.mean(vox_by_scans, axis=0)
rms = np.sqrt(scan_means)
plt.plot(rms)

rms

assert True
assert 1 == 1

def difference_rms(img_arr):
    """ Your code goes here """
    # You want to return the RMS from this function

result = difference_rms(data)
# assert np.all(result == rms) # This should work

worst = np.argmax(rms)

worst

worst_vol = difference_volumes[:, :, :, worst]
plt.gray()
plt.imshow(worst_vol[:, :, 17], interpolation='nearest')

def montage(vol3d):
    """Returns a 2d image montage as an array, given a 3d volume array

    The code here is a little fancy, no need to worry about the details for now.  Just pretend it's FSL
    """
    n_slices = vol3d.shape[-1]
    n_cols = int(np.ceil(np.sqrt(n_slices)))
    # Split and distrubute the 3D array over the last (slice) axis
    rows = np.array_split(vol3d, n_cols, axis=2)
    # ensure the last row is the same size as the others
    rows[-1] = np.dstack((rows[-1], np.zeros(rows[-1].shape[0:2] + (rows[0].shape[2]-rows[-1].shape[2],))))
    im = np.vstack([np.squeeze(np.hstack(np.dsplit(row, n_cols))) for row in rows])
    return(im)

plt.imshow(montage(worst_vol), interpolation='nearest')

in_bad_order = np.argsort(rms)
in_bad_order

def replace_vol(img_arr, vol_no):
    """ Replace volume number `vol_no` with mean of vols either side """
    # Copy the original data, ``img_arr``
    # Take the mean of volumes 64, 66
    # Replace volume 65 with the mean of 64, 66
    
fixed = replace_vol(data, 65) # Call our new function
assert not np.all(fixed == data) # We changed the array
left = data[:, :, :, 64]
right = data[:, :, :, 66]
mean_either_side = (left + right) / 2.0
# assert np.all(fixed[:, :, :, 65] == mean_either_side) # This should work

import awesome # Yes, I wrote it
fixed = awesome.replace_vol(data, 65)
new_rms = awesome.difference_rms(fixed)
plt.plot(rms)
plt.plot(new_rms)