%pylab inline
import numpy as np
import matplotlib.pyplot as plt

import nibabel as nib

img = nib.load('srabold.nii.gz')

print(img.header)

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

data = img.get_data()
data.shape

plt.imshow(data.mean(axis=-1)[:, :, 30])

event_types = ['house', 'scrambled', 'cat', 'shoe', 'bottle', 'scissors', 'chair', 'face']

import os
HOME = os.path.expanduser('~')
DATA_PATH = os.path.join(HOME, 'data', 'ds105')
SUB1_TASK_PATH = os.path.join(DATA_PATH, 'sub001', 'model', 'model001', 'onsets', 'task001_run001')
os.path.isdir(SUB1_TASK_PATH)

SUB1_TASK_PATH

open?

os.path.isdir(SUB1_TASK_PATH)

os.listdir(SUB1_TASK_PATH)

fobj = open(COND1_FNAME, 'rt')
lines = fobj.readlines()
for line in open(COND1_FNAME, 'rt'):
    pass
file.__iter__

COND1_FNAME = os.path.join(SUB1_TASK_PATH, 'cond001.txt')
print(open(COND1_FNAME, 'rt').read())

def evdefs_from_fname(fname):
    """ Get onset, duration, amplitude for event from file
    
    Parameters
    ----------
    fname : str
        Filename to read data from
        
    Returns
    -------
    evdefs : list of lists
        list, one element per event, where the elements are lists containing
        (onset, duration, amplitude)
    """
    evdefs = []
    for line in open(fname, 'rt'):
        values = line.split()
        float_values = [float(v) for v in values]
        evdefs.append(float_values)
    return evdefs

# Of course we need to test
evdefs = evdefs_from_fname(COND1_FNAME)
assert len(evdefs) == 12
assert evdefs[0] == [156.000, 0.5, 1]
assert evdefs[-1] == [178.000, 0.5, 1]

dtype_def = [('onset', 'f'), ('duration', 'f'), ('amplitude', 'f')]
evs = np.loadtxt(COND1_FNAME, dtype_def)
evs

evdefs[2][2]

evs[0]['onset']

evs['onset']

event_defs = {}
for eno in range(8):
    fname = os.path.join(os.path.join(SUB1_TASK_PATH, 'cond%03d.txt' % (eno + 1)))
    print(fname)
    event_name = event_types[eno]
    event_defs[event_name] = np.loadtxt(fname, dtype_def)
event_defs          

import nipy.modalities.fmri.design as fmrid

fmrid.block_design?

nipy_dtype = [('start', 'f'), ('end', 'f'), ('type', 'S10')]
n_events = sum([len(evs) for e in event_defs])
nipy_block_specs = np.zeros((n_events,), dtype=nipy_dtype)
event_no = 0
for ev_type in event_defs:
    ev_def = event_defs[ev_type]
    for e in ev_def:
        nipy_block_specs[event_no]['start'] = e['onset']
        nipy_block_specs[event_no]['end'] = (e['onset'] + e['duration'])
        nipy_block_specs[event_no]['type'] = ev_type
        event_no += 1
nipy_block_defs = np.sort(nipy_block_specs)
nipy_block_specs

n_trs = img.shape[-1]
TR = 2.5
t = np.arange(n_trs) * TR
X_exper, cons_exper = fmrid.block_design(nipy_block_defs, t, level_contrasts=True)
plt.figure(figsize=(10, 20))
plt.imshow(X_exper)

np.set_printoptions(suppress=True)
cons_exper

drift = fmrid.natural_spline(t, [t.mean()])
# Make confounds a little easier to display
drift = drift / drift.max(axis=0)
plt.figure(figsize=(10, 20))
plt.imshow(drift)

plt.plot(drift[:, 3])

X, cons = fmrid.stack_designs((X_exper, cons_exper),
                              (drift, {}))
plt.figure(figsize=(20, 10))
plt.imshow(X)

cons

# Reshape the data for estimation
slice_size = np.prod(data.shape[:-1])
Y = np.reshape(data, (slice_size, n_trs)).T
Y.shape

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

B.shape

cons

contrast = cons['constant_0']
contrast

con_data = contrast.dot(B)
con_data.shape

con_img = np.reshape(con_data, img.shape[:-1])
con_img.shape

plt.imshow(con_img[..., 30])