#!/home/a.ghaderi/.conda/envs/envjm/bin/python
# Model 2
import pystan 
import pandas as pd
import numpy as np
import sys
sys.path.append('../../')
import utils

parts = 1 
data = utils.get_data()          #loading dateset
data = data[data['participant']==parts]

mis = np.where((data['n200lat']<.101)|(data['n200lat']>.248))[0] # missing data for n200lat
obs = np.where((data['n200lat']>.101)&(data['n200lat']<.248))[0] # observation and missing data for n200lat

N_mis = mis.shape[0]  # number of missing data
N_obs = obs.shape[0]   # number of observed data

modelfile = '../../stans/res_nonhier.stan' #reading the model span
f = open(modelfile, 'r')
model_wiener = f.read()
sm = pystan.StanModel(model_code=model_wiener)# Compile the model stan

ncohers = 2  #Number of coherence conditions
nspats = 2   #Number of spatial conditions
nconds = 4   #Number of conditions
y = data['y'].to_numpy()
cond_coher = data['cond_coher'].to_numpy()
cond_spat = data['cond_spat'].to_numpy()
conds = data['conds'].to_numpy()
n200lat = data['n200lat'].to_numpy()

#set inistial data for molde span
data_winner =  {'N_obs':N_obs,      #Number of trial-level observations
                'N_mis':N_mis,      #Number of trial-level mising data
                'ncohers':ncohers,  #Number of coherence conditions
                'nspats':nspats,    #Number of spatial conditions
                'nconds':nconds,    #Number of conditions
                'y':np.concatenate([y[obs],y[mis]]),     #acc*rt in seconds for obervation and missing data
                'cond_coher':np.concatenate([cond_coher[obs],cond_coher[mis]]),     #Coherence index for each trial 
                'cond_spat':np.concatenate([cond_spat[obs],cond_spat[mis]]),        #sptial index for each trial 
                'conds':np.concatenate([conds[obs],conds[mis]]),        #sptial index for each trial
                'n200lat_obs':n200lat[obs]};    #n200 latency for each trial observation

# setting MCMC arguments
niter = 10000
nwarmup = 4000
nchains = 1
thin = 1

initials = [] # initial sampling
for c in range(0, nchains):
    chaininit = {
        'delta': np.random.uniform(1, 3, size=ncohers),
        'alpha': np.random.uniform(.5, 1.),
        'eta': np.random.uniform(.01, .2),
        'res': np.random.uniform(.01, .02, size=nspats),     
        'n200sub': np.random.uniform(.11, .2, size=nconds),
        'lambda': np.random.uniform(.01, .02),
        'n200lat_mis': np.random.uniform(.11, .2, size = N_mis)
    }
    initials.append(chaininit)    
      
# Train the model and generate samples
fit = sm.sampling(data=data_winner, iter=niter, chains=nchains, warmup=nwarmup, thin=thin, init=initials)

utils.to_pickle(stan_model=sm, stan_fit=fit, save_path='../../save/nonhier/'+str(parts)+'_res_nonhier.pkl')


[((7, 0, 7, 25), 'sys.path.append', 'sys.path.append', ({(7, 16, 7, 24): '"""../../"""'}, {}), "('../../')", False, 'import sys\n'), ((11, 7, 11, 23), 'utils.get_data', 'utils.get_data', ({}, {}), '()', False, 'import utils\n'), ((23, 5, 23, 46), 'pystan.StanModel', 'pystan.StanModel', (), '', False, 'import pystan\n'), ((14, 6, 14, 61), 'numpy.where', 'np.where', ({(14, 15, 14, 60): "((data['n200lat'] < 0.101) | (data['n200lat'] > 0.248))"}, {}), "((data['n200lat'] < 0.101) | (data['n200lat'] > 0.248))", True, 'import numpy as np\n'), ((15, 6, 15, 61), 'numpy.where', 'np.where', ({(15, 15, 15, 60): "((data['n200lat'] > 0.101) & (data['n200lat'] < 0.248))"}, {}), "((data['n200lat'] > 0.101) & (data['n200lat'] < 0.248))", True, 'import numpy as np\n'), ((40, 20, 40, 51), 'numpy.concatenate', 'np.concatenate', ({(40, 35, 40, 50): '[y[obs], y[mis]]'}, {}), '([y[obs], y[mis]])', True, 'import numpy as np\n'), ((41, 29, 41, 78), 'numpy.concatenate', 'np.concatenate', ({(41, 44, 41, 77): '[cond_coher[obs], cond_coher[mis]]'}, {}), '([cond_coher[obs], cond_coher[mis]])', True, 'import numpy as np\n'), ((42, 28, 42, 75), 'numpy.concatenate', 'np.concatenate', ({(42, 43, 42, 74): '[cond_spat[obs], cond_spat[mis]]'}, {}), '([cond_spat[obs], cond_spat[mis]])', True, 'import numpy as np\n'), ((43, 24, 43, 63), 'numpy.concatenate', 'np.concatenate', ({(43, 39, 43, 62): '[conds[obs], conds[mis]]'}, {}), '([conds[obs], conds[mis]])', True, 'import numpy as np\n'), ((55, 17, 55, 54), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n'), ((56, 17, 56, 42), 'numpy.random.uniform', 'np.random.uniform', ({(56, 35, 56, 37): '(0.5)', (56, 39, 56, 41): '(1.0)'}, {}), '(0.5, 1.0)', True, 'import numpy as np\n'), ((57, 15, 57, 41), 'numpy.random.uniform', 'np.random.uniform', ({(57, 33, 57, 36): '(0.01)', (57, 38, 57, 40): '(0.2)'}, {}), '(0.01, 0.2)', True, 'import numpy as np\n'), ((58, 15, 58, 55), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n'), ((59, 19, 59, 58), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n'), ((60, 18, 60, 45), 'numpy.random.uniform', 'np.random.uniform', ({(60, 36, 60, 39): '(0.01)', (60, 41, 60, 44): '(0.02)'}, {}), '(0.01, 0.02)', True, 'import numpy as np\n'), ((61, 23, 61, 63), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n')]