%pylab inline

import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['image.cmap'] = 'gray'
plt.rcParams['image.interpolation'] = 'nearest'

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,))))
data = img.get_data()
slice_size = np.prod(data.shape[:-1])
print(slice_size)
Y = np.reshape(data, (slice_size, n_vols)).T
Y.shape

X

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

B.shape

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

B.shape

plt.imshow(B[0, :, :, 40])

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

Yselected = 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')

Yhat = X.dot(Bselected)
plt.plot(Yhat, 'r')
plt.plot(Yselected)

def fit_model(data, X):
    # Shape of design
    n_vols, n_params = X.shape
    assert n_vols == data.shape[-1]
    # Reshape data to be time by voxels
    Y = np.reshape(data, (-1, n_vols)).T
    B = np.linalg.pinv(X).dot(Y)
    # Reshape betas to image
    B = np.reshape(B, (n_params,) + data.shape[:-1])
    return B

# What should we add to X?

B_better = fit_model(data, X_better)

B_better_selected = B_better[:, 19, 12, 40]
Yhat_better = X_better.dot(B_better_selected)
plt.plot(Yhat_better, 'r')
plt.plot(Yselected)

B_better.shape

plt.imshow(B_better[0, :, :, 40])

import nipy.modalities.fmri.hrf as hrfs

t = np.linspace(0, 30, 100)
plt.plot(t, hrfs.glovert(t))
plt.xlabel('time (seconds)')