# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

# + [markdown] deletable=true editable=true
# # Step 5
#
# ## Predict behavioral performance from the match of the gradients
#
# We combine steps 2 and 4 to investigate a potential behavioral relevance of the match of time series to gradients

# + deletable=true editable=true
# %matplotlib inline

import numpy as np
import h5py as h5
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from scipy.stats import spearmanr

from fgrad.predict import features_targets, predict_performance

# + [markdown] deletable=true editable=true
# ## Prepare the features
#
# Same as in Step 3, we simply calculate the average values for each 100-subject group

# + deletable=true editable=true
f = h5.File('/Users/marcel/projects/HCP/volumes_embedded_full.hdf5')
d_LR = f['Working_memory/Run1']
d_RL = f['Working_memory/Run2']

# + deletable=true editable=true
labels = dict()
labels['WM_fix'] = 0
labels['WM_0back'] = 1
labels['WM_2back'] = 2

# Block onsets expressed as TRs
# We add 6 volumes (4.32 s) to each onset to take into account hemodynamic lag 
# and additional 4 volumes (2.88 s) to account for instruction
nback_LR_2b = np.round(np.array([7.977, 79.369, 150.553, 178.689])/0.72).astype(int)+10
nback_LR_0b = np.round(np.array([36.159, 107.464, 221.965, 250.18])/0.72).astype(int)+10
nback_RL_2b = np.round(np.array([7.977, 79.369, 178.769, 250.22])/0.72).astype(int)+10
nback_RL_0b = np.round(np.array([36.159, 107.464, 150.567, 222.031])/0.72).astype(int)+10
nback_fix = np.array([88, 187, 286])+6

# Each block lasts for 27.5 seconds
vols_2b_LR = np.concatenate([range(x,x+38) for x in nback_LR_2b])
vols_0b_LR = np.concatenate([range(x,x+38) for x in nback_LR_0b])
vols_2b_RL = np.concatenate([range(x,x+38) for x in nback_RL_2b])
vols_0b_RL = np.concatenate([range(x,x+38) for x in nback_RL_0b])
vols_fix = np.concatenate([range(x,x+22) for x in nback_fix])
vols_fix = np.concatenate([vols_fix, range(395, 405)])

# Targets
nback_targets_LR = np.zeros(405)
nback_targets_LR[vols_2b_LR] = 1
nback_targets_LR[vols_fix] = -1

nback_targets_RL = np.zeros(405)
nback_targets_RL[vols_2b_RL] = 1
nback_targets_RL[vols_fix] = -1

# + deletable=true editable=true
# Get random group assignments
subjects = f['Working_memory/Subjects'][...]
np.random.seed(123)
sind = np.arange(len(subjects))
G1 = sorted(np.random.choice(sind, 100, replace = False ))
sind = np.delete(sind,G1)
G2 = sorted(np.random.choice(sind, 100, replace = False ))
sind = np.delete(sind,G2)
G3 = sorted(np.random.choice(sind, 100, replace = False ))

# + deletable=true editable=true
conds = ['WM_fix', 'WM_0back', 'WM_2back']
grads = [0,1,2]

# Group 1

f_WM_train1, t_WM_train1 = features_targets(data = d_LR, subjects = G1, 
                                      inds = nback_targets_LR, condnames = conds, gradients = grads, labels = labels)

f_WM_train2, t_WM_train2 = features_targets(data = d_RL, subjects = G1, 
                                      inds = nback_targets_RL, condnames = conds, gradients = grads, labels = labels)

# Group 2

f_WM_train3, t_WM_train3 = features_targets(data = d_LR, subjects = G2, 
                                      inds = nback_targets_LR, condnames = conds, gradients = grads, labels = labels)

f_WM_train4, t_WM_train4 = features_targets(data = d_RL, subjects = G2, 
                                      inds = nback_targets_RL, condnames = conds, gradients = grads, labels = labels)

# Group 3

f_WM_test1, t_WM_test1 = features_targets(data = d_LR, subjects = G3, 
                                      inds = nback_targets_LR, condnames = conds, gradients = grads, labels = labels)

f_WM_test2, t_WM_test2 = features_targets(data = d_RL, subjects = G3, 
                                      inds = nback_targets_RL, condnames = conds, gradients = grads, labels = labels)

# + [markdown] deletable=true editable=true
# ## Prepare targets
#
# Get the estimated d-prime and bias and try to predict them from the gradient match

# + deletable=true editable=true
data_performance = pd.read_csv('../data/WM_SDT.csv', index_col=0)

data_performance.subject_id = data_performance.subject_id.astype('str')

data_performance = data_performance.reset_index(drop = True)

# + deletable=true editable=true
G1_ID = subjects[G1]
G2_ID = subjects[G2]
G3_ID = subjects[G3]

# + deletable=true editable=true
# Indices to include from the behavioral data
WM_ind_G1 = []
WM_ind_G2 = []
WM_ind_G3 = []

# These are the indices of features to remove because of missing data
remove_G1 = []
remove_G2 = []
remove_G3 = []

for i, n in enumerate(G1_ID):
    try:
        WM_ind_G1.append(np.where(data_performance.iloc[:,0] == n)[0][0])
    except:
        remove_G1.append(i)

for i, n in enumerate(G2_ID):
    try:
        WM_ind_G2.append(np.where(data_performance.iloc[:,0] == n)[0][0])
    except:
        remove_G2.append(i)
    
for i, n in enumerate(G3_ID):
    try:
        WM_ind_G3.append(np.where(data_performance.iloc[:,0] == n)[0][0])
    except:
        remove_G3.append(i)

# + [markdown] deletable=true editable=true
# #### Prepare the vectors

# + deletable=true editable=true
t_WM_2b_train1 = data_performance['dprime_2b_LR'][WM_ind_G1]
t_WM_2b_train2 = data_performance['dprime_2b_RL'][WM_ind_G1]
t_WM_2b_train3 = data_performance['dprime_2b_LR'][WM_ind_G2]
t_WM_2b_train4 = data_performance['dprime_2b_RL'][WM_ind_G2]
t_WM_2b_test1 = data_performance['dprime_2b_LR'][WM_ind_G3]
t_WM_2b_test2 = data_performance['dprime_2b_RL'][WM_ind_G3]

t_WM_0b_train1 = data_performance['dprime_0b_LR'][WM_ind_G1]
t_WM_0b_train2 = data_performance['dprime_0b_RL'][WM_ind_G1]
t_WM_0b_train3 = data_performance['dprime_0b_LR'][WM_ind_G2]
t_WM_0b_train4 = data_performance['dprime_0b_RL'][WM_ind_G2]
t_WM_0b_test1 = data_performance['dprime_0b_LR'][WM_ind_G3]
t_WM_0b_test2 = data_performance['dprime_0b_RL'][WM_ind_G3]

f_WM_2b_t1 = np.delete(f_WM_train1[2::3], remove_G1, axis = 0)
f_WM_2b_t2 = np.delete(f_WM_train2[2::3], remove_G1, axis = 0)
f_WM_2b_t3 = np.delete(f_WM_train3[2::3], remove_G2, axis = 0)
f_WM_2b_t4 = np.delete(f_WM_train4[2::3], remove_G2, axis = 0)
f_WM_2b_test1 = np.delete(f_WM_test1[2::3], remove_G3, axis = 0)
f_WM_2b_test2 = np.delete(f_WM_test2[2::3], remove_G3, axis = 0)

f_WM_0b_t1 = np.delete(f_WM_train1[1::3], remove_G1, axis = 0)
f_WM_0b_t2 = np.delete(f_WM_train2[1::3], remove_G1, axis = 0)
f_WM_0b_t3 = np.delete(f_WM_train3[1::3], remove_G2, axis = 0)
f_WM_0b_t4 = np.delete(f_WM_train4[1::3], remove_G2, axis = 0)
f_WM_0b_test1 = np.delete(f_WM_test1[1::3], remove_G3, axis = 0)
f_WM_0b_test2 = np.delete(f_WM_test2[1::3], remove_G3, axis = 0)

# + [markdown] deletable=true editable=true
# ## Predict performance

# + deletable=true editable=true
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import explained_variance_score
import statsmodels.api as sm

# + [markdown] deletable=true editable=true
# ### 2-back

# + deletable=true editable=true
features_A_2b = np.vstack([f_WM_2b_t1, f_WM_2b_t2])
features_B_2b = np.vstack([f_WM_2b_t3, f_WM_2b_t4])
features_C_2b = np.vstack([f_WM_2b_test1, f_WM_2b_test2])

targets_A_2b = np.concatenate([t_WM_2b_train1, t_WM_2b_train2])
targets_B_2b = np.concatenate([t_WM_2b_train3, t_WM_2b_train4])
targets_C_2b = np.concatenate([t_WM_2b_test1, t_WM_2b_test2])

# + deletable=true editable=true
predict_performance(features_A_2b, targets_A_2b, features_B_2b, targets_B_2b, features_C_2b, targets_C_2b)

# + deletable=true editable=true
sns.set_style('whitegrid')
size = 6

fig = plt.figure(figsize = (15,12))

ax = fig.add_subplot(3,3,1)
ax.scatter(features_A_2b[:,0], targets_A_2b, s = size)
ax.set_ylabel('Data split 1', fontsize = 18)
ax.set_title('Gradient 1', fontsize = 18)

ax = fig.add_subplot(3,3,2)
ax.scatter(features_A_2b[:,1], targets_A_2b, s = size)
ax.set_title('Gradient 2', fontsize = 18)

ax = fig.add_subplot(3,3,3)
ax.scatter(features_A_2b[:,2], targets_A_2b, s = size)
ax.set_title('Gradient 3', fontsize = 18)

ax = fig.add_subplot(3,3,4)
ax.scatter(features_B_2b[:,0], targets_B_2b, s = size)
ax.set_ylabel('Data split 2', fontsize = 18)

ax = fig.add_subplot(3,3,5)
ax.scatter(features_B_2b[:,1], targets_B_2b, s = size)

ax = fig.add_subplot(3,3,6)
ax.scatter(features_B_2b[:,2], targets_B_2b, s = size)

ax = fig.add_subplot(3,3,7)
ax.scatter(features_C_2b[:,0], targets_C_2b, s = size)
ax.set_ylabel('Data split 3', fontsize = 18)

ax = fig.add_subplot(3,3,8)
ax.scatter(features_C_2b[:,1], targets_C_2b, s = size)

ax = fig.add_subplot(3,3,9)
ax.scatter(features_C_2b[:,2], targets_C_2b, s = size)

plt.tight_layout()

# + deletable=true editable=true
scaler = StandardScaler()
scaler.fit(features_A_2b[:,[1,2]])
ols_2b = sm.OLS(targets_A_2b, sm.add_constant(scaler.transform(features_A_2b[:,[1,2]])))
ols_2b = ols_2b.fit(cov_type="HC1")
print ols_2b.summary()
pred_2b_B_A = ols_2b.predict(sm.add_constant(scaler.transform(features_A_2b[:,[1,2]])))
pred_2b_B_C = ols_2b.predict(sm.add_constant(scaler.transform(features_C_2b[:,[1,2]])))

print "Variance explained (A): %.2f" % explained_variance_score(targets_A_2b, pred_2b_B_A)
print "Variance explained (C): %.2f" % explained_variance_score(targets_C_2b, pred_2b_B_C)

# + [markdown] deletable=true editable=true
# ### 0-back

# + deletable=true editable=true
scaler = StandardScaler()
scaler.fit(f_WM_0b_t1)

features_A_0b = np.vstack([f_WM_0b_t1, f_WM_0b_t2])
features_B_0b = np.vstack([f_WM_0b_t3, f_WM_0b_t4])
features_C_0b = np.vstack([f_WM_0b_test1, f_WM_0b_test2])

targets_A_0b = np.concatenate([t_WM_0b_train1, t_WM_0b_train2])
targets_B_0b = np.concatenate([t_WM_0b_train3, t_WM_0b_train4])
targets_C_0b = np.concatenate([t_WM_0b_test1, t_WM_0b_test2])

# + deletable=true editable=true
predict_performance(features_A_0b, targets_A_0b, features_B_0b, targets_B_0b, features_C_0b, targets_C_0b)

# + deletable=true editable=true
scaler = StandardScaler()
scaler.fit(features_A_0b[:,2])
ols_0b = sm.OLS(targets_A_0b, sm.add_constant(scaler.transform(features_A_0b[:,2])))
ols_0b = ols_0b.fit(cov_type="HC1")
print ols_0b.summary()
pred_0b_B_A = ols_0b.predict(sm.add_constant(scaler.transform(features_A_0b[:,2])))
pred_0b_B_C = ols_0b.predict(sm.add_constant(scaler.transform(features_C_0b[:,2])))

print "Variance explained (A): %.2f" % explained_variance_score(targets_A_0b, pred_0b_B_A)
print "Variance explained (C): %.2f" % explained_variance_score(targets_C_0b, pred_0b_B_C)

# + deletable=true editable=true
sns.set_style('whitegrid')
size = 6

fig = plt.figure(figsize = (15,12))

ax = fig.add_subplot(3,3,1)
ax.scatter(features_A_0b[:,0], targets_A_0b, s = size)
ax.set_ylabel('Data split 1', fontsize = 18)
ax.set_title('Gradient 1', fontsize = 18)

ax = fig.add_subplot(3,3,2)
ax.scatter(features_A_0b[:,1], targets_A_0b, s = size)
ax.set_title('Gradient 2', fontsize = 18)

ax = fig.add_subplot(3,3,3)
ax.scatter(features_A_0b[:,2], targets_A_0b, s = size)
ax.set_title('Gradient 3', fontsize = 18)

ax = fig.add_subplot(3,3,4)
ax.scatter(features_B_0b[:,0], targets_B_0b, s = size)
ax.set_ylabel('Data split 2', fontsize = 18)

ax = fig.add_subplot(3,3,5)
ax.scatter(features_B_0b[:,1], targets_B_0b, s = size)

ax = fig.add_subplot(3,3,6)
ax.scatter(features_B_0b[:,2], targets_B_0b, s = size)

ax = fig.add_subplot(3,3,7)
ax.scatter(features_C_0b[:,0], targets_C_0b, s = size)
ax.set_ylabel('Data split 3', fontsize = 18)

ax = fig.add_subplot(3,3,8)
ax.scatter(features_C_0b[:,1], targets_C_0b, s = size)

ax = fig.add_subplot(3,3,9)
ax.scatter(features_C_0b[:,2], targets_C_0b, s = size)

plt.tight_layout()

# + [markdown] deletable=true editable=true
# ## Plot relationships for winning models

# + deletable=true editable=true
sns.set_style('white')

dotsize = 10

fig = plt.figure(figsize = (10,10))


ax = plt.subplot2grid((2,2),(1, 0))
ax.scatter(features_B_2b[:,1], targets_B_2b, s = dotsize)
ax.set_xlabel('Gradient 2', fontsize = 16)
ax.set_ylabel('2-back D-Prime', fontsize = 16)
#ax.set_title('2-back', fontsize = 16)

ax = plt.subplot2grid((2,2),(1, 1))
ax.scatter(features_B_2b[:,2], targets_B_2b, s = dotsize)
ax.set_xlabel('Gradient 3', fontsize = 16)
ax.set_ylabel('2-back D-Prime', fontsize = 16)
#ax.set_title('2-back', fontsize = 16)

ax = plt.subplot2grid((2,2),(0, 0), colspan = 2)
ax.scatter(features_B_0b[:,2], targets_B_0b, s = dotsize)
ax.set_xlabel('Gradient 3', fontsize = 16)
ax.set_ylabel('0-back D-Prime', fontsize = 16)
#ax.set_title('2-back', fontsize = 16)

fig.savefig('../figures/grad-performance_B.pdf')

StarcoderdataJupyterScriptsDedupFiltered