%pylab inline

import numpy as np
import matplotlib.pyplot as plt

import os

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

import glob

run_fnames = glob.glob(os.path.join(DATA_PATH, 'sub001', 'BOLD', 'task*', 'bold.nii.gz'))
run_fnames.sort()
run_fnames

from nipy.algorithms.registration import FmriRealign4d
from nipy import load_image, save_image

N_SLICES = 40
TR = 2.5
SLICE_AXIS = 0
space_to_time = list(range(0, N_SLICES, 2)) + list(range(1, N_SLICES, 2))
time_to_space = np.argsort(space_to_time)

# Let's use only the first 2 runs for now
run_fnames = run_fnames[:2]

# Make filenames into images
run_imgs = [load_image(fname) for fname in run_fnames]

R = FmriRealign4d(run_imgs, slice_order=time_to_space, slice_info=(0, 1))
# Estimate motion within- and between- sessions
R.estimate(loops=1, refscan=None)

resample_run0 = R.resample(0)

transforms0 = R._transforms[0]

motion_trans_rotvec = np.vstack([t.param * t.precond[:6] for t in transforms0])

from nipy.externals.transforms3d import taitbryan
from nipy.algorithms.registration.affine import rotation_vec2mat

motion_trans_euler = motion_trans_rotvec.copy()
for scan_no in range(motion_trans_rotvec.shape[0]):
    rot_mat = rotation_vec2mat(motion_trans_rotvec[scan_no, 3:])
    euler_angles = taitbryan.mat2euler(rot_mat)
    motion_trans_euler[scan_no, 3:] = euler_angles

fig, axes = plt.subplots(2, 1)
axes[0].plot(motion_trans_rotvec[:, :3])
axes[0].set_ylabel('translation (mm)')
axes[1].plot(motion_trans_rotvec[:, 3:])
axes[1].set_xlabel('time (TR)')
axes[1].set_ylabel('rotation (radians)')

before_var = np.std(run_imgs[0].get_data(), axis=-1)
after_var = np.std(resample_run0.get_data(), axis=-1)
fig = figure()
fig.set_size_inches(10.5,18.5)
plt.imshow(np.hstack((before_var[..., 40], after_var[..., 40])), cmap="gray", interpolation='nearest')

from subprocess import check_call

check_call('fsl4.1-mcflirt -in %s -plots' % run_fnames[0], shell=True)

pth, fname = os.path.split(run_fnames[0])
par_fname = os.path.join(pth, 'bold_mcf.par')
params = np.loadtxt(par_fname)

fix, axes = plt.subplots(2, 1)
axes[0].plot(params[:, :3])
axes[1].plot(params[:, 3:])

mcsampled_run0 = load_image(os.path.join(pth, 'bold_mcf.nii.gz'))

after_var = np.std(mcsampled_run0.get_data(), axis=-1)
fig = figure()
fig.set_size_inches(10.5,18.5)
plt.imshow(np.hstack((before_var[..., 40], after_var[..., 40])), cmap="gray", interpolation='nearest')