%pylab inline

import numpy as np
import matplotlib.pyplot as plt

import nibabel as nib

import os
HOME = os.path.expanduser('~')
DATA_PATH = os.path.join(HOME, 'data')

img_fname = os.path.join(DATA_PATH, 'ds105', 'sub001', 'BOLD', 'task001_run001', 'bold.nii.gz')
img = nib.load(img_fname)
n_vols = img.shape[-1]
n_vols

# from sub001/model/model001/onsets/task001_run001
block_starts_seconds = np.array([12, 48, 84, 120, 156, 192, 228, 264])
block_length_seconds = 22
TR = 2.5

block_starts_scans = block_starts_seconds / TR
block_length_scans = block_length_seconds / TR
block_starts_scans

delay = 2
regressor = np.zeros((n_vols,))
for start in block_starts_scans:
    start_tr = round(start + delay)
    end_tr = round(start + delay + block_length_scans)
    regressor[start_tr:end_tr+1] = 1
    print start_tr, end_tr
plt.plot(regressor)

X = np.column_stack((regressor, np.ones((n_vols,))))
Y = img.get_data()
slice_size = np.prod(Y.shape[:-1])
print slice_size
Y = np.reshape(Y, (slice_size, n_vols)).T
Y.shape

X

B = np.linalg.pinv(X).dot(Y)
Pinv = np.linalg.pinv(X)
B = Pinv.dot(Y)
B = np.dot(Pinv, Y)

B.shape

B = np.reshape(B, (2,) + img.shape[:-1])

B.shape

plt.imshow(B[0, :, :, 40], cmap='gray')

B_selected = B[:, 19, 12, 40]
B_selected

Yselected = img.get_data()[19, 12, 40]
Yselected.shape

Bselected = np.linalg.pinv(X).dot(Yselected)

Bselected

plt.plot(Yselected - np.mean(Yselected))
plt.plot(regressor * 40, 'r')

Xp = np.linalg.pinv(X)

Xp_again = np.linalg.inv(X.T.dot(X)).dot(X.T)

np.allclose(Xp - Xp_again, 0)

X2 = np.column_stack((X, X.dot([0.5, 0.5])))

# print X2

np.linalg.inv(X2.T.dot(X2))

pinvX2  =  np.linalg.pinv(X2)
print pinvX2.shape