# This analysis requires installation of HDDM. See Wiecki et al., 2013, Front. NeuroInform for more details.
# Current link to software download: http://ski.clps.brown.edu/hddm_docs/

# import files
import pandas as pd
import matplotlib
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
import hddm
# For regressor models
from patsy import dmatrix


# load csv file with single trial data
# for off DBS
off_data = hddm.load_csv('SingleTrialData_off.csv')
# for cDBS (with off DBS)
cDBS_data = hddm.load_csv('SingleTrialData_offcDBS.csv')
# for aDBS
aDBS_data = hddm.load_csv('SingleTrialData_aDBS.csv')
# Check data, e.g.
off_data.head(10)

-----------------------------------------------------
#Models

# off
m_off = hddm.HDDMRegressor(off_data, {"a ~ C(Instruction)+C(Difficulty)", "v ~ C(Difficulty)"})
# for aDBS
m_aDBS = hddm.HDDMRegressor(aDBS_data, {"a ~ C(Instruction)*C(Cuelockedstim)+C(Difficulty)*C(Cuelockedstim)", "v ~ C(Difficulty)*C(Cuelockedstim)", "t ~ C(Cuelockedstim)"})
m_aDBS_posthoc = hddm.HDDMRegressor(aDBS_data, {"a ~ C(Instruction)+C(Difficulty):C(Cuelockedstim)", "v ~ C(Difficulty)"})
# cDBS (test for coherence*stim IA as in m_aDBS)
m_cDBS = hddm.HDDMRegressor(cDBS_data, {"a ~ C(Instruction)+C(Difficulty)*C(stim)", "v ~ C(Difficulty)"})

# for beta ~ threshold regression, e.g
m_betacue = hddm.HDDMRegressor(off_data, {"a ~ C(Instruction)+C(Difficulty)+Betacue", "v ~ C(Difficulty)"})
m_betapost = hddm.HDDMRegressor(off_data, {"a ~ C(Instruction)+C(Difficulty)+Betapost:C(Difficulty)", "v ~ C(Difficulty)"})

# For running the models
m_off.find_starting_values()
m_off.sample(10000, burn=2000) # same for other models

-------------------------------------------------------
# Check results

# Visualize thresholds and drift rates for different conditions (have to be adjusted to actual parameters)

# e.g. for off:
a_Instr = m_off.nodes_db.node["a_C(Instruction)[T.1]"]
a_Coh = m_off.nodes_db.node["a_C(Difficulty)[T.1]"]
v_Coh = m_off.nodes_db.node["v_C(Difficulty)[T.1]"]
# e.g. for aDBS
a_InstrstimIA, a_CohstimIA = m_aDBS.nodes_db.node[["a_C(Instruction)[T.1]:C(Cuelockedstim)[T.1]", "a_C(Difficulty)[T.1]:C(Cuelockedstim)[T.1]"]]
a_Coh_nostim_posthoc, a_Coh_stim_posthoc = m_aDBS_posthoc.nodes_db.node[["a_C(Difficulty)[T.1]:C(Cuelockedstim)[0]", "a_C(Difficulty)[T.1]:C(Cuelockedstim)[1]"]]

# Plot the parameters

# e.g. for off:
hddm.analyze.plot_posterior_nodes([a_Instr], 10)
hddm.analyze.plot_posterior_nodes([a_Coh], 10)
hddm.analyze.plot_posterior_nodes([v_Coh], 10)

# for saving plots, e.g.
hddm.analyze.plot_posterior_nodes([a_Instr], 10);
plt.savefig('a_Instr_off.eps')

# Compare the parameters statistically
# e.g. for off
print "P(Speed < Accuracy) = ", (a_Instr.trace() < 0).mean()
# or aDBS
print "P(Coh_stim < Coh_nostim) = ", (a_Coh_stim.trace() < a_Coh_nostim.trace()).mean()
print "P(Coh_stim > 0) = ", (a_Coh_stim_posthoc.trace() > 0).mean()
print "P(Coh_nostim > 0) = ", (a_Coh_nostim_posthoc.trace() > 0).mean()
--------------------------------------------------------------------
# Model checks, e.g. for off

m_off.plot_posteriors()

models = []
for i in range(5):
  m_off_check = hddm.HDDMRegressor(off_data, {"a ~ C(Instruction)+C(Difficulty)", "v ~ C(Difficulty)"})
  m_off_check.find_starting_values()
  m_off_check.sample(10000, burn=2000)
  models.append(m_off_check)
hddm.analyze.gelman_rubin(models)