# import packages we will need
%pylab inline

# nicer plots, but only available in recent matplotlib
try:
    plt.style.use('ggplot')
except:
    pass

import hrf_estimation as he
print('You are running hrf_estimation version %s' % he.__version__)

voxels, conditions, onsets = he.data.get_sample_data(0, full_brain=False)
print('Num of timepoints in a single voxel: %s' % voxels.shape[0])
print('Num of voxels: %s' % voxels.shape[1])
print('Num of conditions: %s' % len(conditions))
print('Num of onsets: %s' % len(onsets))

### perform detrending: Savitzky-Golay filter
n_voxels = voxels.shape[-1]
voxels = voxels - he.savitzky_golay.savgol_filter(voxels, 91, 3, axis=0)
voxels /= voxels.std(axis=0)
print("Performed detrending of the BOLD timecourse using a Savitzky-Golay filter")

# plot one voxel
plt.figure(figsize=(18, 6))
plt.plot(voxels[:2 * 672, 0], lw=2)
plt.xlabel('Timecourse (in seconds)')
plt.show()

# we construct the matrix of drifts as the matrix of ones that account for
# the intercept term. If you have drifts estimated by e.g. SPM you can add
# them here.
drifts = np.ones((voxels.shape[0], 1))

# we call he.glm with our data and a TR of 1 second. This will return the 
# estimated HRF and the 
# we also specify the basis function for the HRF estimation
hrfs, betas = he.glm(conditions, onsets, 1., voxels, drifts=drifts, mode='r1glm', basis='3hrf', verbose=1)

xx = np.linspace(0, 25) # range of values for time
# construct the final HRF by multiplying by its basis
generated_hrfs = hrfs[0] * he.hrf.spmt(xx)[:, None] + hrfs[1] * he.hrf.dspmt(xx)[:, None] + hrfs[2] * he.hrf.ddspmt(xx)[:, None]

plt.figure(figsize=(12, 6))
plt.plot(xx, generated_hrfs)
plt.ylim((-.5, 1.))
plt.show()

hrfs, betas = he.glm(conditions, onsets, 1., voxels, drifts=drifts, basis='2hrf', verbose=1)
xx = np.linspace(0, 25) # range of values for time
# construct the final HRF by multiplying by its basis
generated_hrfs = hrfs[0] * he.hrf.spmt(xx)[:, None] + hrfs[1] * he.hrf.dspmt(xx)[:, None]

plt.figure(figsize=(12, 6))
plt.plot(xx, generated_hrfs)
plt.ylim((-.5, 1.))
plt.show()

hrfs, betas = he.glm(conditions, onsets, 1., voxels, drifts=drifts, basis='fir', verbose=1)
xx = np.linspace(0, 25) # range of values for time
# construct the final HRF by multiplying by its basis
# generated_hrfs = hrfs[0] * he.hrf.spmt(xx)[:, None] + hrfs[1] * he.hrf.dspmt(xx)[:, None]

plt.figure(figsize=(12, 6))
plt.plot(np.arange(hrfs.shape[0]), hrfs)
plt.ylim((-.5, 1.))
plt.axis('tight')
plt.show()

design_all_sessions = []
voxels_all_sessions = []
drifts_all_sessions = []
n_basis = 20
for sess in range(5):
    voxels, conditions, onsets = he.data.get_sample_data(sess, full_brain=False)

    # detrending
    n_voxels = voxels.shape[-1]
    voxels = voxels - he.savitzky_golay.savgol_filter(voxels, 91, 3, axis=0)
    voxels /= voxels.std(axis=0)
    voxels_all_sessions.append(voxels)

    # create design matrix
    design, _ = he.utils.create_design_matrix(conditions, onsets, 1., voxels.shape[0], hrf_length=n_basis, basis='fir')
    design_all_sessions.append(design)
    drifts_all_sessions.append(np.ones((voxels.shape[0], 1)))

from scipy import sparse
design = sparse.block_diag(design_all_sessions)
drifts = sparse.block_diag(drifts_all_sessions)
out = he.rank_one(design, np.reshape(voxels_all_sessions, (-1, 500)), n_basis, drifts=drifts)

# plot the result
hrfs, beta = out
sign = np.sign(np.dot(hrfs.T, he.hrf.spmt(np.arange(n_basis))))
plt.figure(figsize=(12, 6))
plt.plot(np.arange(hrfs.shape[0]), (sign * hrfs) / np.abs(hrfs).max(0))
#plt.ylim((-.5, 1.))
plt.axis('tight')
plt.show()

hrfs, betas = he.glm(conditions, onsets, 1., voxels, drifts=drifts, mode='glms', basis='3hrf', verbose=1)

hrfs_mean = hrfs.mean(1)
generated_hrfs = hrfs_mean[0] * he.hrf.spmt(xx)[:, None] + hrfs_mean[1] * he.hrf.dspmt(xx)[:, None] + \
    hrfs_mean[2] * he.hrf.ddspmt(xx)[:, None]
sign = np.sign(np.dot(generated_hrfs.T, he.hrf.spmt(xx)))
norm = np.abs(generated_hrfs).max(0)
generated_hrfs = generated_hrfs * sign / norm

plt.figure(figsize=(12, 6))
plt.plot(generated_hrfs)
plt.ylim((-1., 1.))
plt.show()