import pyddm as ddm
from ddm import Sample
from ddm import plot

import os, glob
import copy
import numpy as np
import scipy as sp
from scipy import stats
import pandas as pd
import statsmodels.api as sm
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
from joblib import Parallel, delayed

# from mcginley_pipeline import behavior
import analysis_tools

from IPython import embed as shell

matplotlib.rcParams['pdf.fonttype'] = 42
matplotlib.rcParams['ps.fonttype'] = 42
sns.set(style='ticks', font='Arial', font_scale=1, rc={
    'axes.labelsize': 7,
    'axes.titlesize': 7,
    'xtick.labelsize': 6,
    'ytick.labelsize': 6,
    'legend.fontsize': 6,
    'axes.linewidth': 0.25,
    'xtick.major.width': 0.25,
    'ytick.major.width': 0.25,
    'ytick.major.width': 0.25,
    'ytick.major.width': 0.25,
    'ytick.major.pad' : 2.0,
    'ytick.minor.pad' : 2.0,
    'xtick.major.pad' : 2.0,
    'xtick.minor.pad' : 2.0,
    'axes.labelpad' : 4.0,
    'axes.titlepad' : 6.0,
    } )
sns.plotting_context()
sns.set_palette("tab10")


def params_melt(params, model_settings, flip_b=False):
    try:
        params = params.loc[:,params.columns!='Unnamed: 0']
    except:
        pass
    
    # # flip b & z:
    # if flip_b:
    #     params_overall = pd.DataFrame({'subj_idx': np.array(df_emp.groupby(['subj_idx']).first().index.get_level_values('subj_idx')),
    #                     'c': np.array(df_emp.groupby(['subj_idx']).apply(analyses_tools.behavior, 'c'))})
    #     b_columns = params.columns[[p[0]=='b' for p in params.columns]]
    #     for b_column in b_columns:
    #         params['{}'.format(b_column)] = params[b_column]
    #         for subj in params['subj_idx'].unique():
    #             if params_overall.loc[params_overall['subj_idx'] == subj, 'c'].values < 0:
    #                 params.loc[params['subj_idx'] == subj, '{}'.format(b_column)] = params.loc[params['subj_idx'] == subj, '{}'.format(b_column)] * -1
    #     z_columns = params.columns[[p[0]=='z' for p in params.columns]]
    #     for z_column in z_columns:
    #         params['{}'.format(z_column)] = params[z_column]
    #         for subj in params['subj_idx'].unique():
    #             if params_overall.loc[params_overall['subj_idx'] == subj, 'c'].values < 0:
    #                 params.loc[params['subj_idx'] == subj, '{}'.format(z_column)] = params.loc[params['subj_idx'] == subj, '{}'.format(z_column)] * -1
    
    # melt:
    params = params.melt(id_vars=['subject_id'])
    for i in range(params.shape[0]):
        # variable = "".join(itertools.takewhile((str.isalpha or str=="_"), params.loc[i,'variable']))

        variable = ''.join([l for l in params.loc[i,'variable'] if not l.isnumeric()])

        if variable in model_settings['depends_on']: 
            conditions = model_settings['depends_on'][variable]
            if conditions is not None:
                if len(conditions) == 2:
                    params.loc[i,conditions[0]] = int(params.loc[i,'variable'][-3])
                    params.loc[i,conditions[1]] = int(params.loc[i,'variable'][-1])
                elif len(conditions) == 1:
                    params.loc[i,conditions[0]] = int(params.loc[i,'variable'][-1])
        params.loc[i,'variable'] = variable
    
    return params

def plot_ddm_fits_reward(df, df_sim, fig=None, log=False):

    plt_nr = 1
    if fig is None:
        fig = plt.figure(figsize=(3, 1.75))
    
    for r in [0,1]:
        ax = fig.add_subplot(1,4,plt_nr)
        bins = 100
        # if log:
        #     old = np.seterr(invalid='ignore')
        #     plt.hist(x=np.log(df.loc[(df['reward']==r)&(df['hit']==1), 'rt']-df.loc[(df['reward']==r)&(df['hit']==1), 'start']), 
        #                     bins=bins, density=True, alpha=0.75, linewidth=0, color=sns.color_palette()[r%2], edgecolor='black', histtype='stepfilled')
            
        #     df_sim.loc[(df_sim['reward']==r)&(df_sim['hit']==1), 'rt2'] = np.log(df_sim.loc[(df_sim['reward']==r)&(df_sim['hit']==1), 'rt'] - df_sim.loc[(df_sim['reward']==r)&(df_sim['hit']==1), 'start'])
        #     sns.kdeplot(data=df_sim.loc[(df_sim['reward']==r)&(df_sim['hit']==1),:], x='rt2',
        #                     bw_adjust=0.75, cut=0, color='black', alpha=1, **{'linewidth': 1})

        # else:
        plt.hist(x=df.loc[(df['reward']==r)&(df['hit']==1), 'rt']-df.loc[(df['reward']==r)&(df['hit']==1), 'start'], 
                        bins=bins, density=True, alpha=0.75, linewidth=0, color=sns.color_palette()[r%2], edgecolor='black', histtype='stepfilled')
        
        df_sim.loc[(df_sim['reward']==r)&(df_sim['hit']==1), 'rt2'] = df_sim.loc[(df_sim['reward']==r)&(df_sim['hit']==1), 'rt'] - df_sim.loc[(df_sim['reward']==r)&(df_sim['hit']==1), 'start']
        sns.kdeplot(data=df_sim.loc[(df_sim['reward']==r)&(df_sim['hit']==1),:], x='rt2',
                        bw_adjust=0.75, cut=0, color='black', alpha=1, **{'linewidth': 1})
        plt.xlim(0,3)
        # plt.title('HIT')
        plt.ylabel('')
        plt.yticks([])
        if r == 0:
            plt.xlabel('RT correct lick(s. wrt signal start)')
        ax.tick_params(axis='both', which='major', labelsize=6)
        ax.tick_params(axis='both', which='minor', labelsize=5)

        plt_nr += 1
    
    # plt_nr = 2
    for r in [0,1]:
        ax = fig.add_subplot(1,4,plt_nr)
        start = 14
        bins = 100
        if log:
            old = np.seterr(invalid='ignore')
            plt.hist(x=np.log10(df.loc[(df['reward']==r)&(df['fa']==1), 'rt']), 
                            bins=bins, density=True, alpha=0.75, linewidth=0, color=sns.color_palette()[r%2], edgecolor='black', histtype='stepfilled')
            # plt.hist(x=df_sim.loc[(df_sim['fa']==1), 'rt'], 
            #                 bins=bins, density=True, alpha=0.75, linewidth=1, color='black', edgecolor='black', histtype='step')
            df_sim['rt'] = np.log10(df_sim['rt'])
            print(df_sim['rt'])
            sns.kdeplot(data = df_sim.loc[(df_sim['reward']==r)&(df_sim['fa']==1), :], x='rt', 
                            bw_adjust=0.75, cut=0, color='black', alpha=1, **{'linewidth': 1})
        else:
            plt.hist(x=df.loc[(df['reward']==r)&(df['fa']==1), 'rt'], 
                            bins=bins, density=True, alpha=0.75, linewidth=0, color=sns.color_palette()[r%2], edgecolor='black', histtype='stepfilled')
            # plt.hist(x=df_sim.loc[(df_sim['fa']==1), 'rt'], 
            #                 bins=bins, density=True, alpha=0.75, linewidth=1, color='black', edgecolor='black', histtype='step')
            sns.kdeplot(data = df_sim.loc[(df_sim['reward']==r)&(df_sim['fa']==1), :], x='rt', 
                            bw_adjust=0.75, cut=0, color='black', alpha=1, **{'linewidth': 1})
            plt.xlim(0,11)
        # plt.title('FA')
        if r == 0:
            plt.xlabel('RT incorrect lick (s. wrt trial start)')
        plt.ylabel('')
        plt.yticks([])
        
        ax.tick_params(axis='both', which='major', labelsize=6)
        ax.tick_params(axis='both', which='minor', labelsize=5)
        
        plt_nr += 1
    
    sns.despine(trim=False)
    plt.tight_layout()
    # plt.subplots_adjust(wspace=0.25)
    
    return fig

def plot_ddm_fits_blocks(df, df_sim, fig=None):

    plt_nr = 1
    if fig is None:
        fig = plt.figure(figsize=(4, 2.5))
    
    # ax = fig.add_subplot(1,2,plt_nr)
    
    # f_fa = df['fa'].mean()
    # f_fa_sim = df_sim['fa'].mean()
    # plt.scatter([-1],[f_fa], s=25, marker='o', color=color)
    # plt.scatter([-1],[f_fa_sim], marker='x', color='black')
    
    # # plt.ylabel('FA (% of trials)')
    # plt.xlabel('Noise duration (s)')
    # plt.ylim(0,1)
    # ax = ax.twinx()
    # x = df.groupby(['start'])['hit'].mean().reset_index()['start']
    # y = df.groupby(['start'])['hit'].mean().reset_index()['hit']
    # plt.errorbar(x[x!=14], y[x!=14], fmt='o', color=color)
    # x = df_sim.groupby(['start'])['hit'].mean().reset_index()['start']
    # y = df_sim.groupby(['start'])['hit'].mean().reset_index()['hit']
    # plt.errorbar(x[x!=14], y[x!=14], fmt='-', color='black')
    # # plt.ylabel('H (% of signals)')
    # plt.xticks(np.arange(-1,12,1), ['X'] + list(np.arange(0,12,1)))
    # # plt.scatter([0,1,2], [f_fa, f_hit, f_miss], marker='o', s=20)
    # # plt.scatter([0,1,2], [f_fa_sim, f_hit_sim, f_miss_sim], marker='x', color='black')
    # # plt.xticks([0,1,2], ['FA', 'H', 'M'])
    # plt.ylim(0,1)
    # plt.axvline(-0.5, color='k', lw=0.5)
    # plt_nr+= 1

    for b in [1,2,3,4,5,6]:
        ax = fig.add_subplot(2,6,plt_nr)
        start = 14
        bins = 100
        plt.hist(x=df.loc[(df['block_id']==b)&(df['hit']==1), 'rt']-df.loc[(df['block_id']==b)&(df['hit']==1), 'start'], 
                        bins=bins, density=True, alpha=0.75, linewidth=0, color=sns.color_palette()[b%2], edgecolor='black', histtype='stepfilled')
        # plt.hist(x=df_sim.loc[(df_sim['hit']==1), 'rt']-df_sim.loc[(df_sim['hit']==1), 'start'], 
        #                 bins=bins, density=True, alpha=0.75, linewidth=1, color='black', edgecolor='black', histtype='step')
        
        df_sim.loc[(df_sim['block_id']==b)&(df_sim['hit']==1), 'rt2'] = df_sim.loc[(df_sim['block_id']==b)&(df_sim['hit']==1), 'rt'] - df_sim.loc[(df_sim['block_id']==b)&(df_sim['hit']==1), 'start']
        sns.kdeplot(data=df_sim.loc[(df_sim['block_id']==b)&(df_sim['hit']==1),:], x='rt2',
                        bw_adjust=0.75, cut=0, color='black', alpha=1, **{'linewidth': 1})
        plt.xlim(0,3)
        # plt.title('HIT')
        plt.ylabel('')
        plt.yticks([])
        if b == 3:
            plt.xlabel('RT correct lick (s. wrt signal start)')
        else:
            plt.xlabel('')
        ax.tick_params(axis='both', which='major', labelsize=6)
        ax.tick_params(axis='both', which='minor', labelsize=5)

        plt_nr+= 1
    

    for b in [1,2,3,4,5,6]:
        ax = fig.add_subplot(2,6,plt_nr)
        start = 14
        bins = 100
        plt.hist(x=df.loc[(df['block_id']==b)&(df['fa']==1), 'rt'], 
                        bins=bins, density=True, alpha=0.75, linewidth=0, color=sns.color_palette()[b%2], edgecolor='black', histtype='stepfilled')
        # plt.hist(x=df_sim.loc[(df_sim['fa']==1), 'rt'], 
        #                 bins=bins, density=True, alpha=0.75, linewidth=1, color='black', edgecolor='black', histtype='step')
        sns.kdeplot(data = df_sim.loc[(df_sim['block_id']==b)&(df_sim['fa']==1), :], x='rt', 
                        bw_adjust=0.75, cut=0, color='black', alpha=1, **{'linewidth': 1})
        plt.xlim(0,11)
        # plt.title('FA')
        if b == 3:
            plt.xlabel('RT incorrect lick (s. wrt trial start)')
        else:
            plt.xlabel('')
        plt.ylabel('')
        plt.yticks([])
        
        ax.tick_params(axis='both', which='major', labelsize=6)
        ax.tick_params(axis='both', which='minor', labelsize=5)

        plt_nr+= 1
    
    sns.despine(trim=False)
    plt.tight_layout()
    plt.subplots_adjust(wspace=0.25)
    
    return fig

def plot_bars(params, parameters):

    data = params.melt(id_vars=['subject_id'])
    data = data.loc[np.array([v[:-1] in parameters for v in data['variable']]),:]
    data['reward'] = [v[-1] for v in data['variable']]
    data['variable'] = [v[0] for v in data['variable']]
    # if data.shape[0]==0:
    #     return plt.figure(figsize=(1+(0.75*data.shape[0]/2), 1.75))

    fig = plt.figure(figsize=(0.5+(0.75*data.shape[0]/2), 1.25))
    plt_nr = 1
    for v in parameters:
        if not v[0] in data['variable'].values:
            continue
        ax = fig.add_subplot(1,int(data.shape[0]/2),plt_nr)
        sns.barplot(x='variable', y='value', hue='reward',
                    palette=[sns.color_palette()[0], sns.color_palette()[1],],
                    data=data.loc[data['variable']==v[0],:], lw=0, ax=ax)
        plt.legend([],[], frameon=False)
        plt_nr += 1
    sns.despine(trim=False)
    plt.tight_layout()
    return fig

def plot_trace():

    v = 3
    k = 4
    dt = 0.05
    noise_dur = 5
    tmax = noise_dur + 3
    x = np.arange(0,tmax,dt)
    traces = []
    for _ in range(10000):

        trace = np.zeros(len(x))
        for t in range(1, len(x)):
            if x[t] < noise_dur:
                trace[t] = trace[t-1] - (trace[t-1]*k*dt) + np.random.normal(0,1/np.sqrt(2),1)*np.sqrt(dt)
            else:
                trace[t] = trace[t-1] - (trace[t-1]*k*dt) + (v*dt) + np.random.normal(0,1/np.sqrt(2),1)*np.sqrt(dt)
        traces.append(trace)
    traces = np.vstack(traces)
    from accumodels import plot_tools
    fig = plot_tools.traces_plot(traces, bound=1, dt=0.05, vmax=0.005, ylim=(-0.5,1.5))
    for i in range(10):
        plt.plot(x, traces[i,:], color='k')
    return fig

def add_individual_difference_bins(df):

    df_cox = behavior.cox_prepare_data(df)

    groupby = ['subject_id']
    columns = ['hit', 'd_ps', 'r_ps', 'rt2']

    df_res_high = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'] & (df['reward']==1),:], 
                                                df_cox=df_cox.loc[df_cox['reward_ind'] & (df_cox['reward']==1),:], 
                                                groupby=groupby)
    df_res_high = df_res_high.set_index(groupby)
    df_res_low = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'] & (df['reward']==0),:], 
                                                df_cox=df_cox.loc[df_cox['reward_ind'] & (df_cox['reward']==0),:], 
                                                groupby=groupby)
    df_res_low = df_res_low.set_index(groupby)

    df_res_mean = (df_res_high[columns] + df_res_low[columns])/2
    df_res_diff = (df_res_high[columns] - df_res_low[columns])
    df_res_diff.columns = [c + '_d' for c in df_res_diff.columns]

    df_res = pd.concat((df_res_mean, df_res_diff), axis=1).reset_index()


    df_res['d_bin'] = pd.qcut(df_res['d_ps'], 3, labels=False)
    df_res['c_bin'] = pd.qcut(df_res['r_ps'], 3, labels=False)

    print(df_res.groupby(['d_bin']).mean())
    print(df_res.groupby(['c_bin']).mean())

    # plt.plot(df_res.groupby(['d_bin']).mean()['d_ps'])
    # plt.plot(df_res.groupby(['c_bin']).mean()['d_ps'])

    for subj, d in df.groupby(['subject_id']):
        print(int(df_res.loc[df_res['subject_id']==subj, 'd_bin']))
        df.loc[d.index, 'd_bin'] = int(df_res.loc[df_res['subject_id']==subj, 'd_bin'])
        df.loc[d.index, 'c_bin'] = int(df_res.loc[df_res['subject_id']==subj, 'c_bin'])
    
    return df

def sample_from_condition(model, ids, nr_trials):

    if isinstance(ids, int):
        ids = [ids]

    # sample:
    samp = model.solve({key:value for value, key in zip(ids, model.required_conditions)}).resample(nr_trials)

    # to dataframe:
    df_upper = pd.DataFrame({'rt': samp.choice_upper,})
    df_undecided = pd.DataFrame({'rt': np.repeat(np.NaN, samp.undecided),})
    df = pd.concat((df_upper, df_undecided))
    df['response'] = (~df['rt'].isna()).astype(int)

    # meta:
    for value, key in zip(ids, model.required_conditions):
        df[key] = value
    
    return df

def simulate_data_gonogo(df, params, model_settings, nr_trials=10000, rt_cutoff=0.15, add_conditions_to_sample=False, n_jobs=24):

    import pyddm as ddm
    from ddm import Sample
    from ddm import Fitted
    import pyddm_components
    import analysis_tools

    # params:
    noise_dur_mean = 5
    noise_dur_max = 11
    sig_dur = 3
    
    # normal sessions:
    noise_durs = np.random.exponential(noise_dur_mean, nr_trials)
    noise_durs[noise_durs>=noise_dur_max] = noise_dur_max

    # catch sessions:
    noise_durs_c = np.random.exponential(noise_dur_mean, nr_trials)
    noise_durs_c[noise_durs_c>=noise_dur_max] = noise_dur_max + sig_dur

    # combine:
    noise_durs = np.concatenate((noise_durs[:int((df['catch_session']==0).mean()*nr_trials)],
                                 noise_durs_c[:int((df['catch_session']==1).mean()*nr_trials)]))

    # noise_durs = noise_durs[noise_durs<=noise_dur_max]
    df_sample = pd.DataFrame({'noise_dur': noise_durs, 'duration': np.repeat(sig_dur, len(noise_durs))})
    df_sample['response'] = 1
    df_sample['rt'] = 1

    # bin
    shift = 'mean'
    df_sample['start'] = df_sample['noise_dur'].copy()
    df_sample.loc[(df_sample['start']<11), 'bins'] = pd.cut(df_sample.loc[(df_sample['start']<11), 'start'], bins=[-1,rt_cutoff,0.5,1,1.5,2,2.5,3,3.5,4,5,6,7,8,9,10,11], labels=False)
    df_sample.loc[df_sample['start']==11, 'bins'] = df_sample['bins'].max()+1
    df_sample.loc[df_sample['start']==14, 'bins'] = df_sample['bins'].max()+1
    print(df_sample.groupby(['bins'])['start'].mean())

    if shift == 'min':
        df_sample['start'] = df_sample.groupby(['bins'])['start'].transform(lambda x: x.min()).round(1)
    if shift == 'max':
        df_sample['start'] = df_sample.groupby(['bins'])['start'].transform(lambda x: x.max()).round(1)
    if shift == 'mean':
        df_sample['start'] = df_sample.groupby(['bins'])['start'].transform(lambda x: x.mean()).round(1)
    print(df_sample.groupby(['bins'])['start'].mean())

    # check if needs to be repeated:
    if add_conditions_to_sample:
        conditions = []
        for key, value in model_settings['depends_on'].items():
            if (value is not None):
                for s in value:
                    if (s not in conditions):
                        conditions.append(s)
        
        dfs_sample = []
        for c in conditions:
            for v in df[c].unique():
                d = df_sample.copy()
                d[c] = v
                print(d)
                dfs_sample.append(d)
        df_sample = pd.concat(dfs_sample)

    # make sample:
    sample = Sample.from_pandas_dataframe(df=df_sample, rt_column_name='rt', correct_column_name='response')

    # make model:
    model = pyddm_components.make_model_one_accumulator(sample=sample, model_settings=model_settings)

    # set fitted parameters:
    param_names = model.get_model_parameter_names()
    try:
        params_to_set = [Fitted(params[param]) for param in param_names]
    except Exception as e:
        print(e)
        print(param_names)
    model.set_model_parameters(params_to_set)
    
    # compute number of trials to generate:
    nr_trials = df_sample.shape[0]
    df_sample.loc[:,'trials'] = 1
    trial_nrs = df_sample.groupby(model.required_conditions).count()['trials']
    trial_nrs = (trial_nrs / sum(trial_nrs) * nr_trials).reset_index()
    trial_nrs['trials'] = trial_nrs['trials'].astype(int)
    print(trial_nrs)
    
    # resample:
    res = Parallel(n_jobs=n_jobs)(delayed(sample_from_condition)(model=model, ids=ids, nr_trials=int(d['trials']))
                                    for ids, d in trial_nrs.groupby(model.required_conditions))
    
    # concat:
    df_sim = pd.concat(res)

    # columns:
    df_sim.loc[df_sim['rt']>=14, 'rt'] = np.NaN
    df_sim.loc[df_sim['rt']>=(df_sim['start']+3), 'response'] = 0
    df_sim.loc[df_sim['rt']>=(df_sim['start']+3), 'rt'] = np.NaN
    df_sim['noise_dur'] = df_sim['start']
    # df_sim['reward'] = ((df_sim['block_id']%2)==1).astype(int)
    df_sim['start'] = df_sim['start'].round(1)

    # sdt columns:
    df_sim = analysis_tools.add_sdt_columns(df_sim)

    # too fast:
    df_sim = df_sim.loc[(df_sim['rt']>=rt_cutoff)|(df_sim['rt'].isna())]
    # df_sim = analysis_tools.fast_hits_to_false_alarms(df_sim)

    # df_sim.loc[df_sim['fa']==1, 'start'] = 14
    # df_sim.loc[df_sim['cr']==1, 'start'] = 14

    return df_sim

n_jobs = 6
n_cpus = 64
exp = 'behavior'
rt_cutoff = 0.15

# load data:
project_dir = 'ddm/'
df = pd.read_csv(os.path.join(project_dir, 'data_ddm_{}.csv'.format(exp)))
fig_dir = '/Users/jwdegee/Library/CloudStorage/Box-Box/McGinleyLabFiles/projects/2021_Attentional_Effort/figs'
df['response'] = df['choice'].astype(int)
df['duration'] = 3
df.loc[df['noise_dur'].isna(), 'noise_dur'] = 14

# load:
print(df.groupby(['state']).count())

# assign random rts:
df = analysis_tools.assign_random_rts(df)

# shell()

#################################
# bin data:
#################################

# bin
shift = 'mean'
df['start'] = df['noise_dur'].copy()
df.loc[(df['start']<11), 'bins'] = pd.cut(df.loc[(df['start']<11), 'start'], bins=[-1,rt_cutoff,0.5,1,1.5,2,2.5,3,3.5,4,5,6,7,8,9,10,11], labels=False)
df.loc[df['start']==11, 'bins'] = df['bins'].max()+1
df.loc[df['start']==14, 'bins'] = df['bins'].max()+1
print(df.groupby(['bins'])['start'].mean())

if shift == 'min':
    df['start'] = df.groupby(['bins'])['start'].transform(lambda x: x.min()).round(1)
if shift == 'max':
    df['start'] = df.groupby(['bins'])['start'].transform(lambda x: x.max()).round(1)
if shift == 'mean':
    df['start'] = df.groupby(['bins'])['start'].transform(lambda x: x.mean()).round(1)
print(df.groupby(['bins'])['start'].mean())

# correct:
plt.hist(df.loc[df['outcome']==0, 'rt'], bins=500, alpha=0.5)
print((df['rt']<rt_cutoff).sum())
df.loc[df['outcome']==0, 'rt_corr'] = df.loc[df['outcome']==0, 'rt'] - (df.loc[df['outcome']==0, 'noise_dur']-df.loc[df['outcome']==0, 'start'])
df.loc[(df['outcome']==0)&(df['rt_corr']>rt_cutoff), 'rt'] = df.loc[(df['outcome']==0)&(df['rt_corr']>rt_cutoff), 'rt_corr']
print((df['rt']<=rt_cutoff).sum())
# df.loc[(df['outcome']==0)&(df['rt']<0.1), 'rt'] = 0.1
# plt.hist(df['rt'], bins=500, alpha=0.5)
# plt.hist(df.loc[df['outcome']==0, 'rt_corr'], bins=500, alpha=0.5)
plt.hist(df.loc[df['outcome']==0, 'rt'], bins=500, alpha=0.5)
plt.close('all')

df.loc[df['rt'].isna(), 'rt'] = 100
df['block_id'] = df['block_id'].astype(int)

# only reward_ind trials:
df = df.loc[df['block_id']>0,:]
df = df.loc[df['reward_ind'],:]
# df = df.loc[df['reward']==0,:]

# params:
dx = 0.005
dt = 0.01
T_dur = 14.5

model_0 = {'a': 'fit', 'v':'fit', 't':'fit', 'k':'fit', 'l':'fit', 'z':0, 'sz':0, 'dc':'fit', 's':0, 'm':0,
            'depends_on': {
                        'v':['reward'],  
                        'k':['reward'], 
                        'l':['reward'], 
                        'dc':['reward'], 
                        't':['reward'],
                        'a':None, 'z':None, 'sz':None,  's':None, 'm':None
                        },
            'dx': dx, 'dt': dt, 'T_dur': T_dur}

model_1 = {'a': 'fit', 'v':'fit', 't':'fit', 'k':'fit', 'l':'fit', 'z':0, 'sz':0, 'dc':'fit', 's':0, 'm':0,
            'depends_on': {
                        'v':['reward'], 
                        'a':['reward'], 
                        'l':['reward'], 
                        'dc':['reward'],
                        't':['reward'], 
                        'k':None, 'z':None, 'sz':None,  's':None, 'm':None
                        },
            'dx': dx, 'dt': dt, 'T_dur': T_dur}

model_2 = {'a': 'fit', 'v':'fit', 't':'fit', 'k':'fit', 'l':'fit', 'z':0, 'sz':0, 'dc':'fit', 's':0, 'm':0,
            'depends_on': {
                        'v':['reward'], 
                        'k':['reward'], 
                        'l':['reward'], 
                        'a':['reward'],
                        't':['reward'],
                        'dc':None, 'z':None, 'sz':None,  's':None, 'm':None
                        },
            'dx': dx, 'dt': dt, 'T_dur': T_dur}



model_3 = {'a': 'fit', 'v':'fit', 't':'fit', 'k':'fit', 'l':0, 'z':0, 'sz':0, 'dc':'fit', 's':0, 'm':0,
            'depends_on': {
                        'v':['reward'],  
                        'k':['reward'], 
                        'dc':['reward'], 
                        't':['reward'],
                        'l':None, 'a':None, 'z':None, 'sz':None,  's':None, 'm':None
                        },
            'dx': dx, 'dt': dt, 'T_dur': T_dur}

model_4 = {'a': 'fit', 'v':'fit', 't':'fit', 'k':'fit', 'l':'fit', 'z':0, 'sz':0, 'dc':'fit', 's':0, 'm':0,
            'depends_on': {
                        'v':['reward'],  
                        'k':['reward'], 
                        'dc':['reward'], 
                        't':['reward'],
                        'l':None, 'a':None, 'z':None, 'sz':None,  's':None, 'm':None
                        },
            'dx': dx, 'dt': dt, 'T_dur': T_dur}


model_5 = {'a': 'fit', 'v':'fit', 't':'fit', 'k':'fit', 'l':'fit', 'z':0, 'sz':0, 'dc':0, 's':0, 'm':0,
            'depends_on': {
                        'v':['reward'],  
                        'k':['reward'], 
                        'l':['reward'], 
                        'dc':None, 
                        't':['reward'],
                        'a':None, 'z':None, 'sz':None,  's':None, 'm':None
                        },
            'dx': dx, 'dt': dt, 'T_dur': T_dur}

model_6 = {'a': 'fit', 'v':'fit', 't':'fit', 'k':'fit', 'l':'fit', 'z':0, 'sz':0, 'dc':'fit', 's':0, 'm':0,
            'depends_on': {
                        'v':['reward'],  
                        'k':['reward'], 
                        'l':['reward'], 
                        'dc':None, 
                        't':['reward'],
                        'a':None, 'z':None, 'sz':None,  's':None, 'm':None
                        },
            'dx': dx, 'dt': dt, 'T_dur': T_dur}

model_settings = [model_0, model_1, model_2, model_3, model_4, model_5, model_6]

## PLOT BICS: #################
###############################

# groupby = ['block_id', 'hit_p']
# groupby = ['reward']
groupby = []
bics = []
for version in [0,1,2,3,4,5,6]:
    filenames = glob.glob(os.path.join(project_dir, 'ddm_params_{}_*').format(version))
    filenames = [f for f in filenames if 'sample' not in f]
    if len(filenames) == 0:
        continue
    params = pd.concat([pd.read_csv(f)for f in filenames])
    params = params.loc[params['bic']==params['bic'].min(),:]
    bics.append(params['bic'].mean())
    print(params['bic'])

bics = np.array(bics).ravel()
bics = bics - bics[0]
x = np.arange(len(bics))

fig = plt.figure(figsize=(2,2))
# ax = fig.add_subplot(111)
plt.bar(x, bics)
sns.despine(trim=False)
plt.tight_layout()
fig.savefig(os.path.join(fig_dir, 'bics.pdf'))

## PLOT FITS: #################
###############################

shell()

# for version in [0,1,2,3,4,5,6]:
for version in [0]:

    filenames = glob.glob(os.path.join(project_dir, 'ddm_params_{}_*').format(version))
    filenames = [f for f in filenames if 'sample' not in f]
    if len(filenames) == 0:
        continue
    params = pd.concat([pd.read_csv(f)for f in filenames])
    params = params.loc[params['bic']==params['bic'].min(),:]
    params['subject_id'] = 0
    # params['block_id'] = params['block_id'].astype(int)
    # params['reward'] = (params['block_id']%2).astype(int)
    params['noise'] = 1
    print()
    print(version)
    print(params)

    # if version != 5:
    #     params['a'] = params['a0/a1']

    # shell()

    # plot model fits:
    #####################################################################################################################

    df_sim = simulate_data_gonogo(df=df.loc[:,['start', 'duration', 'response', 'rt', 'reward', 'catch_session']], 
                                                params=params,
                                                model_settings=model_settings[version], 
                                                nr_trials=500000, 
                                                rt_cutoff=rt_cutoff,
                                                add_conditions_to_sample=True,
                                                n_jobs=n_jobs).reset_index()
    df_sim['catch'] = 0
    fig = plot_ddm_fits_reward(df, df_sim, log=False)
    fig.savefig(os.path.join(fig_dir, 'hist_{}.pdf'.format(version)))

    groupby = ['reward']
    df_sim_res = analysis_tools.compute_metrics(df=df_sim, groupby=groupby)
    df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'],:], 
                                                groupby=groupby)
    
    df_res_sem = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'],:], 
                                            groupby=['subject_id', 'reward'])[['reward','tr_c', 'tr_d', 'hit']].groupby(['reward']).sem().reset_index()

    fig = plt.figure(figsize=(1.5*3,1.5))
    # for i, m in enumerate(['tr_hr', 'tr_far', 'hit', 'rt2']):
    for i, m in enumerate(['tr_c', 'tr_d', 'hit']):
        ax = fig.add_subplot(1,3,i+1)

        plt.errorbar(df_res.loc[(df_res['reward']==0),m], 
                    df_sim_res.loc[(df_sim_res['reward']==0),m], 
                    xerr=df_res_sem.loc[(df_res_sem['reward']==0),m], 
                    fmt='x', markersize=10, color=sns.color_palette()[0])
        plt.errorbar(df_res.loc[(df_res['reward']==1),m], 
                    df_sim_res.loc[(df_sim_res['reward']==1),m], 
                    xerr=df_res_sem.loc[(df_res_sem['reward']==1),m], 
                    fmt='x', markersize=10, color=sns.color_palette()[1])
        # plt.scatter(df_res.loc[(df_res['reward']==1),m], 
        #             df_sim_res.loc[(df_sim_res['reward']==1),m], 
        #             marker='o', color=sns.color_palette()[1])
        # plt.scatter(df_res.loc[(df_res['reward']==1),m], 
        #             df_sim_res.loc[(df_sim_res['reward']==1),m], 
        #             marker='x', color=sns.color_palette()[1])
        minimum = min((df_res[m].min(), df_sim_res[m].min()))
        maximum = max((df_res[m].max(), df_sim_res[m].max()))
        difference = maximum-minimum
        plt.plot(np.linspace(minimum-difference, maximum+difference, 10), 
                 np.linspace(minimum-difference, maximum+difference, 10), 
                 color='black', lw=0.5, ls='--')
        plt.xlabel('{} data'.format(m))
        plt.ylabel('{} model'.format(m))
        lim = ax.get_xlim()
        ax.set_xlim(lim)
        ax.set_ylim(lim)
        # ticks = ax.get_xticks()
        # ax.set_xticks(ticks)
        # ax.set_yticks(ticks)
        plt.locator_params(nbins=3)
        # plt.xticks(np.arange(1,7),np.arange(1,7))
    sns.despine(trim=False)
    plt.tight_layout()
    fig.savefig(os.path.join(fig_dir, 'fits_{}.pdf'.format(version)))

    # shell()
    # dfs_sim = []
    # for b in [1,2,3,4,5,6]:
    #     for p in [0,1]:
    #         print(p)
    #         df_sim = simulate_data_gonogo(df=df.loc[(df['block_id']==b)&(df['hit_p']==p),['start', 'duration', 'response', 'rt', 'catch_session']], 
    #                                                     params=params.loc[(params['block_id']==b)&(params['hit_p']==p),:],
    #                                                     model_settings=model_settings[version], 
    #                                                     nr_trials=500000, 
    #                                                     rt_cutoff=rt_cutoff,
    #                                                     n_jobs=n_jobs).reset_index()
    #         df_sim['hit_p'] = p
    #         df_sim['block_id'] = b
    #         df_sim['reward'] = b%2
    #         df_sim['subject_id'] = 0
    #         dfs_sim.append(df_sim)
    # df_sim = pd.concat(dfs_sim)
    # df_sim.to_csv(os.path.join(data_dir, 'data_sim_{}_{}.csv'.format(exp, version)))

    # exclude trial that are too fast:
    # df_sim = df_sim.loc[(df_sim['rt']>=rt_cutoff)|(df_sim['rt'].isna()),:]

    # df_sim = analysis_tools.fast_hits_to_false_alarms(df_sim, rt_cutoff=rt_cutoff)

    # df = df.loc[(df['rt']>=rt_cutoff)|(df['rt'].isna())]

    # shell()

    # # plot:
    # for r in [0,1]:
    #     # for p in [0]:
    #     fig = plot_ddm_fits(df.loc[(df['reward']==r),:], 
    #                         df_sim.loc[(df_sim['reward']==r),:], 
    #                         color=sns.color_palette()[r%2])
    #     fig.savefig(os.path.join(project_dir, 'figs', 'hist_{}_{}_{}.pdf'.format(version, groupby[0], r)))
    
    # for p in [0,1]:
    #     fig = plot_ddm_fits(df.loc[(df['hit_p']==p),:], 
    #                         df_sim.loc[(df_sim['hit_p']==p),:])
    #     fig.savefig(os.path.join(project_dir, 'figs', 'hist_{}_{}_{}_{}.pdf'.format(version, groupby[0], groupby[1], p)))

shell()

## PLOT RESULTS: #################
##################################

for version in [0]:

    model_settings = {'a': 0.995, 'v':'fit', 't':'fit', 'k':'fit', 'l':'fit', 'z':0, 'sz':0, 'dc':'fit', 's':0, 'm':0,
            'depends_on': {
                        'v':None,  
                        'k':None, 
                        'l':None, 
                        'dc':None, 
                        't':None,
                        'a':None, 'z':None, 'sz':None,  's':None, 'm':None
                        },
            'dx': dx, 'dt': dt, 'T_dur': T_dur}
    
    # for groupby in [['block_id', 'hit_p'],]: # ['reward', 'state']
    for groupby in [['block_id', 'hit_p'], 
                    ['reward', 'pupil_trial_start_bin'],
                    ['reward', 'pupil_trial_start_c3_bin'],]:

        # load params:
        params = []
        for i in range(100):

            if len(groupby) == 1:
                filename = os.path.join(project_dir, 'ddm_params_{}_{}_sample_{}.csv'.format(version, groupby[0], i))
            elif len(groupby) == 2:
                filename = os.path.join(project_dir, 'ddm_params_{}_{}_{}_sample_{}.csv'.format(version, groupby[0], groupby[1], i))
            if os.path.exists(filename):
                p = pd.read_csv(filename)
                p['subject_id'] = i+1
                if (p['k']<0).sum() > 0:
                    print(filename)
                params.append(p)
        params = pd.concat(params).reset_index()

        # correlation matrix:
        if groupby == ['block_id', 'hit_p']:
            columns = ['t', 'v', 'l', 'k', 'dc']

            fig = plt.figure(figsize=(6,8))
            plt_nr = 1
            corrs = []
            for (b,h), data in params.groupby(['block_id', 'hit_p']):
                ax = fig.add_subplot(4,3,plt_nr)
                corr = data[columns].corr()
                corrs.append(corr)
                mask = np.triu(np.ones_like(corr, dtype=bool))
                cmap = "vlag"
                sns.heatmap(corr, mask=mask, cmap=cmap, vmin=-1, vmax=1, center=0,
                    square=True, linewidths=.5, cbar_kws={"shrink": .5}, ax=ax)
                plt_nr += 1
            plt.tight_layout()
            fig.savefig(os.path.join(fig_dir, 'fits_{}_{}_{}_correlation_matrices.pdf'.format(version, groupby[0], groupby[1])))

            corr = pd.DataFrame(np.stack(corrs).mean(axis=0))
            corr.columns = columns
            corr.set_index(keys=np.array(columns), inplace=True)
            mask = np.triu(np.ones_like(corr, dtype=bool))
            cmap = "vlag"
            fig = plt.figure(figsize=(2,2))
            ax = fig.add_subplot(111)
            sns.heatmap(corr, mask=mask, cmap=cmap, vmin=-1, vmax=1, center=0,
                square=True, linewidths=.5, cbar_kws={"shrink": .5}, ax=ax)
            plt.tight_layout()
            fig.savefig(os.path.join(fig_dir, 'fits_{}_{}_{}_correlation_matrix.pdf'.format(version, groupby[0], groupby[1])))

        def plot_bootstrap_estimates(params, x_measure, y_measure, ls='-', hue=None, ax=None, x_values=None):
            
            import matplotlib.transforms as transforms

            if x_measure == 'block_id':
                sns.violinplot(data=params, x=x_measure, y=y_measure, hue=hue,
                    split=True, linewidth=1, ax=ax, scale='width', cut=0,
                    palette=[sns.color_palette('tab10')[1], sns.color_palette('tab10')[0], sns.color_palette('tab10')[1], 
                                sns.color_palette('tab10')[0], sns.color_palette('tab10')[1], sns.color_palette('tab10')[0]])
                mean = params.loc[params['reward']==0,:].groupby(['block_id'])[y_measure].mean().reset_index()
                plt.plot([1,3,5], mean[y_measure], color=sns.color_palette('tab10')[0], linewidth=1.5, ls=ls)
                mean = params.loc[params['reward']==1,:].groupby(['block_id'])[y_measure].mean().reset_index()
                plt.plot([0,2,4], mean[y_measure], color=sns.color_palette('tab10')[1], linewidth=1.5, ls=ls)
                p1 = (params.loc[(params['reward']==1), y_measure].values > 
                        params.loc[(params['reward']==0), y_measure].values).mean()
                p2 = (params.loc[(params['block_id']==1)|(params['block_id']==2), y_measure].values > 
                        params.loc[(params['block_id']==5)|(params['block_id']==6), y_measure].values).mean()
                p3 = ((params.loc[(params['block_id']==1), y_measure].values - params.loc[(params['block_id']==2), y_measure].values) > 
                        (params.loc[(params['block_id']==5), y_measure].values - params.loc[(params['block_id']==6), y_measure].values)).mean()
                p1 = min((p1, 1-p1))
                p2 = min((p2, 1-p2))
                p3 = min((p3, 1-p3))
                plt.text(x=0, y=0.8, s='p = {}\np = {}\np = {}'.format(round(p1,3), round(p2,3), round(p3,3)), size=6, 
                            transform=transforms.blended_transform_factory(ax.transData, ax.transAxes))
                plt.xlabel('Block #')
            elif x_measure == 'state_bin':
                # sns.violinplot(data=params, x=x, y=y, hue=hue,
                #     split=True, linewidth=1, ax=ax, scale='width', cut=0,
                #     color='grey')

                fit = False
                if (params['reward']==0).mean() == 1:
                    color = sns.color_palette('tab10')[0]
                elif (params['reward']==1).mean() == 1:
                    color = sns.color_palette('tab10')[1]
                else:
                    color = 'grey'
                    fit = True

                # for (b, p) in params.loc[params['state_bin']!=100,:].groupby(['subject_id']):
                #     plt.plot(x_values[:-1], p[y], color=color, lw=0.3, alpha=0.1)
                # # plt.plot(np.arange(len(p[y])), params.loc[params['state_bin']!=100,:].groupby(['state_bin'])[y].mean(), color=color)

                x = x_values[:-1]
                y = params.loc[(params['state_bin']!=100), :].groupby(['state_bin'])[y_measure].median()
                y_sem = (params.loc[(params['state_bin']!=100), :].groupby(['state_bin'])[y_measure].quantile(0.975) - 
                         params.loc[(params['state_bin']!=100), :].groupby(['state_bin'])[y_measure].quantile(0.025))/2
                plt.errorbar(x, y, yerr=y_sem, fmt='o', color=color, markerfacecolor='lightgrey', markeredgewidth=0.75)

                y_sem = (params.loc[params['state_bin']==100,y_measure].quantile(0.975) - 
                         params.loc[params['state_bin']==100,y_measure].quantile(0.025)) / 2
                plt.errorbar(x_values[-1], params.loc[params['state_bin']==100,y_measure].median(), 
                             yerr=y_sem, fmt='x', color=color)

                if fit:
                    
                    print()
                    # with all subject values included:
                    df_seq_reg = params.loc[(params['state_bin']!=100), ['subject_id', x_measure, y_measure]].groupby(['state_bin'])[[x_measure, y_measure]].mean()
                    # df_seq_reg[x_measure] = df_seq_reg[x_measure]-df_seq_reg[x_measure].mean()
                    F_values, p_values, model_dfs, resid_dfs, model = analysis_tools.sequential_regression(df=df_seq_reg, x_measure=x_measure, y_measure=y_measure, order=2)
                    print()
                    print(df_seq_reg)
                    print(F_values)
                    print(model_dfs)
                    print(resid_dfs)
                    print(p_values)
                    try:
                        popt, pcov = analysis_tools.curve_fit(analysis_tools.arbitrary_poly, x, y, p0=[1]*(model+1))
                        # print(popt) # This contains your three best fit parameters
                        xx = np.linspace(min(x), max(x), 100)
                        curve_linexp = analysis_tools.arbitrary_poly(xx, *popt)
                        plt.plot(xx, curve_linexp, lw=1.5, ls='-', color='grey', zorder=0)
                    except Exception as e:
                        print(e)
                        pass



                # plt.scatter(np.ones(len(params['subject_id'].unique()))*x_values[-1],
                #             params.loc[params['state_bin']==100,y], color=color, alpha=0.1, s=2)
                # plt.scatter(13, params.loc[params['state_bin']==100,y].mean(), color=color, alpha=0.25)

            ax.legend([],[], frameon=False)
            plt.ylabel(y_measure)
            sns.despine(trim=False)
            plt.tight_layout()
            return None

        def compute_weigthed_mean(params, df, columns=['k', 'v', 't', 'dc', 'l', 'bic']):
            
            weigths = df.groupby(['block_id'])['hit_p'].mean()
            for i, b in enumerate([1,2,3,4,5,6]):
                params.loc[(params['block_id']==b)&(params['hit_p']==0), columns] = params.loc[(params['block_id']==b)&(params['hit_p']==0), columns] * (1-weigths.iloc[i])
                params.loc[(params['block_id']==b)&(params['hit_p']==1), columns] = params.loc[(params['block_id']==b)&(params['hit_p']==1), columns] * weigths.iloc[i]
            params = params.groupby(['block_id', 'subject_id']).sum().reset_index()

            return params
        
        if groupby == ['block_id', 'hit_p']:
            params_m = compute_weigthed_mean(params, df)
        elif 'pupil_trial_start' in groupby[1]:
            params['state_bin'] = params[groupby[1]].astype(int)
            params_m = params.groupby(['subject_id', 'state_bin']).mean().reset_index()

        if 'block_id' in params.columns:
            params_m['reward'] = params_m['block_id']%2
            params['reward'] = params['block_id']%2

        # plot point plots:
        #####################################################################################################################

        for y in ['v', 't', 'k', 'l', 'dc']:
        
            if y in params.columns:
                # fig = analysis_tools.plot_scalars_across_blocks(params, x='block_id', y=y, error='conf')
                # fig.savefig(os.path.join(project_dir, 'figs', 'across_blocks_{}_{}_{}.pdf'.format(version, groupby[0], y)))
                if groupby == ['block_id', 'hit_p']:
                    fig = plt.figure(figsize=(1.75,1.75))
                    ax = fig.add_subplot(111)
                    plot_bootstrap_estimates(params_m, x_measure='block_id', y_measure=y, ax=ax)
                    if y == 't':
                        plt.ylim(0,0.2)
                    fig.savefig(os.path.join(fig_dir, 'fits_across_blocks_{}_{}_{}_{}.pdf'.format(version, groupby[0], groupby[1], y)))

                    fig = plt.figure(figsize=(1.75,1.75))
                    ax = fig.add_subplot(111)
                    plot_bootstrap_estimates(params.loc[params['hit_p']==0,:], x_measure='block_id', y_measure=y, ax=ax, ls='--')
                    plot_bootstrap_estimates(params.loc[params['hit_p']==1,:], x_measure='block_id', y_measure=y, ax=ax, ls='-')
                    if y == 't':
                        plt.ylim(0,0.2)
                    fig.savefig(os.path.join(fig_dir, 'fits_across_blocks_hitp_{}_{}_{}_{}.pdf'.format(version, groupby[0], groupby[1], y)))

                elif 'pupil_trial_start' in groupby[1]:
                    

                    pupil_size = df.groupby(['subject_id'] + groupby)['pupil_trial_start_c3'].mean().groupby(groupby).mean()


                    fig = plt.figure(figsize=(1.75,1.75))
                    ax = fig.add_subplot(111)
                    x_values = pupil_size.groupby(groupby[1]).mean().values * 100
                    plot_bootstrap_estimates(params_m, x_measure='state_bin', y_measure=y, ax=ax, x_values=x_values)
                    plt.xlabel('Pupil size (% max)')
                    if y == 't':
                        plt.ylim(0,0.2)
                    fig.savefig(os.path.join(fig_dir, 'fits_across_blocks_{}_{}_{}_{}.pdf'.format(version, groupby[0], groupby[1], y)))

                    fig = plt.figure(figsize=(1.75,1.75))
                    ax = fig.add_subplot(111)
                    x_values = pupil_size.loc[0,:].values * 100
                    plot_bootstrap_estimates(params.loc[params['reward']==0,:], x_measure='state_bin', y_measure=y, ax=ax, ls='--', x_values=x_values)
                    x_values = pupil_size.loc[1,:].values * 100
                    plot_bootstrap_estimates(params.loc[params['reward']==1,:], x_measure='state_bin', y_measure=y, ax=ax, ls='-', x_values=x_values)
                    plt.xlabel('Pupil size (% max)')
                    if y == 't':
                        plt.ylim(0,0.2)
                    fig.savefig(os.path.join(fig_dir, 'fits_across_blocks_reward_{}_{}_{}_{}.pdf'.format(version, groupby[0], groupby[1], y)))



                # elif groupby == ['reward', 'state']:

                #     fig = plt.figure(figsize=(2,1.75))
                #     ax = fig.add_subplot(111)

                #     sns.violinplot(data=params, x="state", y=y,
                #         split=True, linewidth=1, ax=ax, scale='width', cut=0,
                #         palette=['lightskyblue', 'limegreen', 'tomato', 'lightgray'])
                #     mean = params.groupby(['state'])[y].mean().reset_index()
                #     plt.plot([0,1,2], mean[y].iloc[:-1], color=sns.color_palette('tab10')[0], linewidth=1.5)
                    

                #     # sns.violinplot(data=params, x="state", y=y, hue='reward',
                #     #     split=True, linewidth=1, ax=ax, scale='width', cut=0)

                #     ax.legend([],[], frameon=False)
                #     # if y == 'k':
                #     #     ax.set_yscale('log')
                #     plt.ylim(ylim)
                #     p1 = (params.loc[(params['state']==1), y].values > 
                #             params.loc[(params['state']==0), y].values).mean()
                #     p2 = (params.loc[(params['state']==2), y].values > 
                #             params.loc[(params['state']==1), y].values).mean()
                #     p3 = (params.loc[(params['state']==3), y].values > 
                #             params.loc[(params['state']==2), y].values).mean()
                #     p1 = min((p1, 1-p1))
                #     p2 = min((p2, 1-p2))
                #     p3 = min((p3, 1-p3))
                #     plt.text(x=0, y=0.8, s='p = {}\np = {}\np = {}'.format(round(p1,3), round(p2,3), round(p3,3)), size=6, 
                #                 transform=transforms.blended_transform_factory(ax.transData, ax.transAxes))
                #     plt.xlabel('State')
                #     plt.ylabel(y)
                #     sns.despine(trim=False)
                #     plt.tight_layout()
                #     fig.savefig(os.path.join(project_dir, 'figs', 'across_blocks_{}_{}_{}.pdf'.format(version, groupby[1], y)))



    # shell()
    # # df_sim.loc[df['rt']<df['start'], 'start'] = 14

    # print(df['fa'].mean())
    # print(df_sim['fa'].mean())

    # # df['block_id'] = df['block_id']+1
    # # df_sim['block_id'] = df_sim['block_id']+1
    # # params['block_id'] = params['block_id']+1

    # try:
    #     for b in [1,2,3,4,5,6]:
    #         fig = plot_ddm_fits(df.loc[df['block_id']==b,:], df_sim.loc[df_sim['block_id']==b,:], color=sns.color_palette()[b%2])
    #         fig.savefig(os.path.join(project_dir, 'figs', 'hist_{}_{}_{}.pdf'.format(version, 'block_id', b)))
    # except:
    #     pass

    # for r in [0,1]:
    #     fig = plot_ddm_fits(df.loc[df['reward']==r,:], df_sim.loc[df_sim['reward']==r,:], color=sns.color_palette()[r%2])
    #     fig.savefig(os.path.join(project_dir, 'figs', 'hist_{}_{}_{}.pdf'.format(version, 'reward', r)))
    
    # # # reward in same plot:
    # # r = 0
    # # fig = plot_ddm_fits(df.loc[df['reward']==r,:], df_sim.loc[df_sim['reward']==r,:], color=sns.color_palette()[r%2], fig=None)
    # # r = 1
    # # fig = plot_ddm_fits(df.loc[df['reward']==r,:], df_sim.loc[df_sim['reward']==r,:], color=sns.color_palette()[r%2], fig=fig)
    # # fig.savefig(os.path.join(project_dir, 'figs', 'hist_{}_{}_{}.pdf'.format(version, 'reward', 'combined')))

shell()

# plot bars: #########################
######################################















for version in [0,1,2,3,4,5,6]:

    filenames = glob.glob(os.path.join(project_dir, 'ddm_params_{}_*').format(version))
    filenames = [f for f in filenames if 'sample' not in f]
    if len(filenames) == 0:
        continue
    params = pd.concat([pd.read_csv(f)for f in filenames])
    params = params.loc[params['bic']==params['bic'].min(),:]
    params['subject_id'] = 0
    # params['block_id'] = params['block_id'].astype(int)
    # params['reward'] = (params['block_id']%2).astype(int)
    params['noise'] = 1
    print()
    print(version)
    print(params)

    fig = plot_bars(params, ['a', 'k', 'dc', 'v', 'l', 't'])
    fig.savefig(os.path.join(fig_dir, 'bars_{}.pdf'.format(version)))



    fig = plot_bars(params, ['l', 'z', 'm'])
    fig.savefig(os.path.join(project_dir, 'figs', 'bars_{}_2.pdf'.format(version)))

    fig = plot_bars(params, ['t'])
    fig.savefig(os.path.join(project_dir, 'figs', 'bars_{}_3.pdf'.format(version)))

shell()


    # version = 4
    # for y in ['a', 'v', 't', 'k', 'l', 'm']:

    #     fig = plt.figure(figsize=(2,1.75))
    #     ax = fig.add_subplot(111)

    #     for groupby in [['reward', 'd_bin'], ]:

    #         # load params:
    #         params = []
    #         for i in range(100):
    #             filename = os.path.join(project_dir, 'ddm_params_{}_{}_{}_sample_{}.csv'.format(version, groupby[0], groupby[1], i))
    #             if os.path.exists(filename):
    #                 p = pd.read_csv(filename)
    #                 p['subject_id'] = i+1
    #                 params.append(p)
    #         params = pd.concat(params).reset_index()
    #         # params = params.groupby(['subject_id', groupby[1]]).mean().reset_index()
    #         # print(params.groupby([groupby[0], groupby[1]]).mean())

    #         sns.pointplot(x=groupby[1], y=y, hue=groupby[0], data=params, linewidth=1, ax=ax)
    #         # mean = params.groupby(['state'])[y].mean().reset_index()
    #         # plt.plot([0,1,2], mean[y].iloc[:-1], color=sns.color_palette('tab10')[0], linewidth=1.5)
            

    #         # sns.violinplot(data=params, x="state", y=y, hue='reward',
    #         #     split=True, linewidth=1, ax=ax, scale='width', cut=0)

    #     ax.legend([],[], frameon=False)
    #     # if y == 'k':
    #     #     ax.set_yscale('log')
    #     if y == 'a':
    #         plt.ylim(0.5,1.5)
    #     if y == 't':
    #         plt.ylim(0,0.2)
    #     if y == 'mixture':
    #         plt.ylim(0,0.3)
    #     # p1 = (params.loc[(params['state']==1), y].values > 
    #     #         params.loc[(params['state']==0), y].values).mean()
    #     # p2 = (params.loc[(params['state']==2), y].values > 
    #     #         params.loc[(params['state']==1), y].values).mean()
    #     # p3 = (params.loc[(params['state']==3), y].values > 
    #     #         params.loc[(params['state']==2), y].values).mean()
    #     # p1 = min((p1, 1-p1))
    #     # p2 = min((p2, 1-p2))
    #     # p3 = min((p3, 1-p3))
    #     # plt.text(x=0, y=0.8, s='p = {}\np = {}\np = {}'.format(round(p1,3), round(p2,3), round(p3,3)), size=6, 
    #     #             transform=transforms.blended_transform_factory(ax.transData, ax.transAxes))
    #     plt.xlabel(groupby[1])
    #     plt.ylabel(y)
    #     sns.despine(trim=False)
    #     plt.tight_layout()
    #     fig.savefig(os.path.join(project_dir, 'figs', 'across_blocks_{}_{}_{}.pdf'.format(version, groupby[1], y)))


    # shell()

    # # plot model traces:
    # #####################################################################################################################

    # fig = plot_trace()
    # fig.savefig(os.path.join(project_dir, 'figs', 'traces.pdf'))

    #########################################################
    # plot optimal leak:
    #########################################################

#     def simulate_in_parallel(df, params, model_settings, k=None, v=None, lapse=None, a=None):

#         # from accumodels import pyddm_tools

#         if k is not None:
#             params['k'] = k
#         if a is not None:
#             params['a'] = a
#             params['a0'] = a
#             params['a1'] = a
#             params['a0/a1'] = a
#         if v is not None:
#             params['v'] = v
#         if lapse is not None:
#             params['lapse'] = lapse
#         df_sim = simulate_data_gonogo(df=df, 
#                                         params=params, 
#                                         model_settings=model_settings, 
#                                         nr_trials=500000, 
#                                         n_jobs=1)
#         df_sim['k'] = float(params['k'])
#         # df_sim['a'] = float(params['a'])
#         df_sim['v'] = float(params['v'])
#         df_sim['l'] = float(params['l'])
#         df_sim['reward'] = int(params['reward'])
#         return df_sim


#     if version == 5:
        

#         shell()

#         # plot optimal leak:
#         ##########################################################################################################################

#         groupby = ['reward']
#         params = params.groupby(groupby).mean().reset_index()
#         globals().update(locals())
#         res = Parallel(n_jobs=64, verbose=1, backend='loky')(delayed(simulate_in_parallel)
#                                                                         (data[['start', 'duration', 'response', 'rt', 'hit_p']], 
#                                                                         params.loc[params[groupby[0]]==g,:], 
#                                                                         model_settings[version], 
#                                                                         k, None, None, None)
#                                                                             for g, data in df.groupby(groupby) 
#                                                                                 for k in np.linspace(0,10,51))
#         df_sim = pd.concat(res)
#         df_sim.loc[df_sim['rt']>=14, 'rt'] = np.NaN
#         df_sim.loc[df_sim['rt']>=(df_sim['start']+3), 'response'] = 0
#         df_sim.loc[df_sim['rt']>=(df_sim['start']+3), 'rt'] = np.NaN
#         df_sim['noise_dur'] = df_sim['start']
#         df_sim['start'] = df_sim['start'].round(1)
#         df_sim = analysis_tools.add_sdt_columns(df_sim)
        
#         df_sim_cox = behavior.cox_prepare_data(df_sim)
#         df_res = analysis_tools.compute_metrics(df=df_sim, df_cox=df_sim_cox, groupby=['reward', 'k'])
#         df_res.to_csv(os.path.join(project_dir, 'simulated_block_id_k.csv'))
        
#         fig = plt.figure(figsize=(3.5,1.75))
#         ax = fig.add_subplot(121)
#         for b, d in df_res.groupby(['reward']):
#             plt.plot(d['k'], d['hit'])
#         plt.axvline(params['k'].iloc[0], color=sns.color_palette()[0])
#         plt.axvline(params['k'].iloc[1], color=sns.color_palette()[1])
#         plt.xlabel('Leak')
#         plt.ylabel('Correct lick-rate (% trials)')
#         ax = fig.add_subplot(122)
#         for b, d in df_res.groupby(['reward']):
#             plt.plot(d['k'], d['d_ps'])
#         plt.axvline(params['k'].iloc[0], color=sns.color_palette()[0])
#         plt.axvline(params['k'].iloc[1], color=sns.color_palette()[1])
#         plt.xlabel('Leak')
#         plt.ylabel('Signal-selective responsiveness')
#         sns.despine(trim=False)
#         plt.tight_layout()
#         fig.savefig(os.path.join(project_dir, 'figs', 'simulated_k_{}_{}_{}.pdf'.format(version, groupby[0], 'hit')))

#         shell()

#         # make landscape 1:
#         ##########################################################################################################################
        
#         globals().update(locals())
#         res = Parallel(n_jobs=64, verbose=1, backend='loky')(delayed(simulate_in_parallel)
#                                                                         (df[['start', 'duration', 'response', 'rt', 'hit_p']], 
#                                                                         params.mean(axis=0), 
#                                                                         model_settings[version], 
#                                                                         k, None, None, a)
#                                                                         for k in np.linspace(0,10,11)
#                                                                             for a in np.linspace(0.5,1.5,11))
#         df_sim = pd.concat(res)
#         df_sim.loc[df_sim['rt']>=14, 'rt'] = np.NaN
#         df_sim.loc[df_sim['rt']>=(df_sim['start']+3), 'response'] = 0
#         df_sim.loc[df_sim['rt']>=(df_sim['start']+3), 'rt'] = np.NaN
#         df_sim['noise_dur'] = df_sim['start']
#         df_sim['start'] = df_sim['start'].round(1)
#         df_sim = analysis_tools.add_sdt_columns(df_sim)
        
#         df_sim_cox = behavior.cox_prepare_data(df_sim)
#         df_res = analysis_tools.compute_metrics(df=df_sim, df_cox=df_sim_cox, groupby=['k', 'a'])
#         df_res.to_csv(os.path.join(project_dir, 'simulated_k_a.csv'))

#         # make landscape 2:
#         ##########################################################################################################################
        
#         globals().update(locals())
#         res = Parallel(n_jobs=64, verbose=1, backend='loky')(delayed(simulate_in_parallel)
#                                                                         (df[['start', 'duration', 'response', 'rt', 'hit_p']], 
#                                                                         params.mean(axis=0), 
#                                                                         model_settings[version], 
#                                                                         None, v, lapse, None)
#                                                                         for v in np.linspace(0,5,11)
#                                                                             for lapse in np.linspace(0.25,.75,11))
#         df_sim = pd.concat(res)
#         df_sim.loc[df_sim['rt']>=14, 'rt'] = np.NaN
#         df_sim.loc[df_sim['rt']>=(df_sim['start']+3), 'response'] = 0
#         df_sim.loc[df_sim['rt']>=(df_sim['start']+3), 'rt'] = np.NaN
#         df_sim['noise_dur'] = df_sim['start']
#         df_sim['start'] = df_sim['start'].round(1)
#         df_sim = analysis_tools.add_sdt_columns(df_sim)
        
#         df_sim_cox = behavior.cox_prepare_data(df_sim)
#         df_res = analysis_tools.compute_metrics(df=df_sim, df_cox=df_sim_cox, groupby=['v', 'lapse'])
#         df_res.to_csv(os.path.join(project_dir, 'simulated_v_lapse.csv'))

#         def plot_landscape(df, x='v', y='a', measure='d', vmin=0, vmax=2, cmap='bwr'):
            
#             m = df[[y,x,measure]].pivot(y,x,measure)
#             m.index = np.round(m.index,5)
#             m.columns = np.round(m.columns,5)
#             f, ax = plt.subplots(figsize=(2,2))
#             # sns.heatmap(m, square=True, linewidths=0, cmap=cmap, vmin=m.to_numpy().ravel().mean()-1, vmax=m.to_numpy().ravel().mean()+1, ax=ax)
#             sns.heatmap(m, square=True, linewidths=0, vmin=vmin, vmax=vmax, cmap=cmap, ax=ax)
#             # sns.heatmap(m, square=True, linewidths=0, cmap=cmap, ax=ax)
#             f.axes[0].invert_yaxis()
#             plt.xlabel(x)
#             plt.ylabel(y)
#             plt.setp(ax.get_xticklabels(), rotation=90, ha="right", rotation_mode="anchor")
#             plt.setp(ax.get_yticklabels(), rotation=0, ha="right", rotation_mode="anchor")
#             plt.tight_layout()
#             return f

#         df_res = pd.read_csv(os.path.join(project_dir, 'simulated_v_lapse.csv'))
#         for y, vmin, vmax in zip(['r_ps', 'd_ps', 'hit', 'rt2'], [0, 0, 10, 0.4], [1, 0.6, 45, 1.6]):
#             fig = plot_landscape(df_res, x='v', y='lapse', measure=y, vmin=vmin, vmax=vmax, cmap='viridis')
#             fig.savefig(os.path.join(project_dir, 'figs', 'landscapes_v_lapse_{}.pdf'.format(y)))
    
#         df_res = pd.read_csv(os.path.join(project_dir, 'simulated_k_a.csv'))
#         for y, vmin, vmax in zip(['r_ps', 'd_ps', 'hit', 'rt2'], [0, 0, 10, 0.4], [1, 0.6, 45, 1.6]):
#             fig = plot_landscape(df_res, x='k', y='a', measure=y, vmin=vmin, vmax=vmax, cmap='viridis')
#             fig.savefig(os.path.join(project_dir, 'figs', 'landscapes_a_k_{}.pdf'.format(y)))

#     # fig = plt.figure(figsize=(3,1.75))
#     # ax = fig.add_subplot(121)
#     # order = ['a', 'v', 't', 'k']
#     # sns.boxplot(x='variable', y='value', hue='reward', order=order, linewidth=0.75, data=params_melted, ax=ax)
#     # sns.stripplot(x='variable', y='value', hue='reward', dodge=True, order=order, jitter=False, linewidth=0.75, data=params_melted, ax=ax)
#     # import matplotlib.transforms as transforms
#     # trans = transforms.blended_transform_factory(ax.transData, ax.transAxes)
#     # # for i, v in enumerate(order):
#     # #     t,p = sp.stats.wilcoxon(params.loc[(params['variable']==v)&(params['reward']==1), 'value'], params.loc[(params['variable']==v)&(params['reward']==0), 'value'])
#     # #     plt.text(x=i, y=0.9, s='p={}'.format(round(p,3)), size=6, transform=trans)
#     # plt.legend([],[], frameon=False)
    
#     # ax = fig.add_subplot(122)
#     # order = ['lapse', 'lapse_slope', 'mixture']
#     # sns.boxplot(x='variable', y='value', hue='reward', order=order, linewidth=0.75, data=params_melted, ax=ax)
#     # sns.stripplot(x='variable', y='value', hue='reward', dodge=True, order=order, jitter=False, linewidth=0.75, data=params_melted, ax=ax)
#     # import matplotlib.transforms as transforms
#     # trans = transforms.blended_transform_factory(ax.transData, ax.transAxes)
#     # # for i, v in enumerate(order):
#     # #     t,p = sp.stats.wilcoxon(params.loc[(params['variable']==v)&(params['reward']==1), 'value'], params.loc[(params['variable']==v)&(params['reward']==0), 'value'])
#     # #     plt.text(x=i, y=0.9, s='p={}'.format(round(p,3)), size=6, transform=trans)
#     # plt.legend([],[], frameon=False)
    
#     # sns.despine(trim=False)
#     # plt.tight_layout()
#     # fig.savefig(os.path.join(project_dir, 'figs', 'bars_{}.pdf'.format(version)))


#     # # plot fits:
#     # bins = 50
#     # for reward in [0,1]:
#     #     plt_nr = 1
#     #     fig = plt.figure(figsize=(8,8))
#     #     for start1, start2 in zip(np.sort(df.loc[df['reward']==reward,'start'].unique()), np.sort(df_sim.loc[df_sim['reward']==reward,'start'].unique())):
#     #         print(start1)
#     #         print(start2)
#     #         ax = fig.add_subplot(4,4,plt_nr)
#     #         plt.hist(x=df.loc[(df['rt']!=-1)&(df['reward']==reward)&(df['start']==start1),'rt']-start1, bins=bins, density=True, alpha=0.75, linewidth=0, edgecolor='black', histtype='stepfilled')
#     #         sns.kdeplot(x=df_sim.loc[(df_sim['reward']==reward)&(df_sim['start']==start2),'rt']-start2, bw_adjust=0.75, cut=0, color='black')
#     #         plt.title(start1)
#     #         plt_nr += 1
#     #     plt.tight_layout()
#     #     fig.savefig(os.path.join(project_dir, 'figs', 'hist2_{}_{}.pdf'.format(version, reward)))


    
#     # # plots parameters:
#     # # params = params.rename({'lapse_slope0': 'slope0', 'lapse_slope1': 'slope1'}, axis=1)
#     # params_melted = params_melt(params, model_settings[version])
#     # # params.loc[params['variable']=='pmixturecoef', 'value'] = params.loc[params['variable']=='pmixturecoef', 'value'] * 10
#     # fig = plt.figure(figsize=(3,1.75))
#     # ax = fig.add_subplot(121)
#     # order = ['a', 'v', 't', 'k']
#     # sns.boxplot(x='variable', y='value', hue='reward', order=order, linewidth=0.75, data=params_melted, ax=ax)
#     # sns.stripplot(x='variable', y='value', hue='reward', dodge=True, order=order, jitter=False, linewidth=0.75, data=params_melted, ax=ax)
#     # import matplotlib.transforms as transforms
#     # trans = transforms.blended_transform_factory(ax.transData, ax.transAxes)
#     # # for i, v in enumerate(order):
#     # #     t,p = sp.stats.wilcoxon(params.loc[(params['variable']==v)&(params['reward']==1), 'value'], params.loc[(params['variable']==v)&(params['reward']==0), 'value'])
#     # #     plt.text(x=i, y=0.9, s='p={}'.format(round(p,3)), size=6, transform=trans)
#     # plt.legend([],[], frameon=False)
    
#     # ax = fig.add_subplot(122)
#     # order = ['lapse', 'lapse_slope', 'mixture']
#     # sns.boxplot(x='variable', y='value', hue='reward', order=order, linewidth=0.75, data=params_melted, ax=ax)
#     # sns.stripplot(x='variable', y='value', hue='reward', dodge=True, order=order, jitter=False, linewidth=0.75, data=params_melted, ax=ax)
#     # import matplotlib.transforms as transforms
#     # trans = transforms.blended_transform_factory(ax.transData, ax.transAxes)
#     # # for i, v in enumerate(order):
#     # #     t,p = sp.stats.wilcoxon(params.loc[(params['variable']==v)&(params['reward']==1), 'value'], params.loc[(params['variable']==v)&(params['reward']==0), 'value'])
#     # #     plt.text(x=i, y=0.9, s='p={}'.format(round(p,3)), size=6, transform=trans)
#     # plt.legend([],[], frameon=False)
    
#     # sns.despine(trim=False)
#     # plt.tight_layout()
#     # fig.savefig(os.path.join(project_dir, 'figs', 'bars_{}.pdf'.format(version)))

#     # # import matplotlib.pyplot as plt
#     # # from scipy.stats import exponnorm
#     # # mu = 3.018066e-01    # [0-0.5]
#     # # sigma = 1.874159e-01 # [0-0.5]
#     # # rate = 3.220722e+01    # [0-100]
#     # # K = 1/(sigma*rate)
#     # # r = exponnorm.rvs(K, loc=mu, scale=sigma, size=10000)
#     # # plt.hist(r, bins=100)
#     # # plt.show()




# # shell()
# # a = sp.stats.gaussian_kde(np.random.exponential(1, int(1e4))).pdf(np.linspace(0,20,100))
# # b = sp.stats.gaussian_kde(np.random.normal(2, 1, int(1e4))).pdf(np.linspace(0,20,100))
# # c = (a*0.2) + (b*0.8)


# # plt.plot(a, label='exp')
# # plt.plot(b, label='normal')
# # plt.plot(c, label='mixed')
# # plt.legend()

# # density = True
# # fig = plt.figure(figsize=(6,3))
# # ax = fig.add_subplot(121)
# # plt.hist(df.loc[(df['rt']!=-1)&(df['reward']==0),'rt'], bins=28, density=density, alpha=0.5)
# # plt.hist(df_sim_0['rt'], bins=28, density=density, alpha=0.5)
# # plt.title('yes: {} ({})\ncorrect: {} ({})'.format(round(yes_0,2), round(yes_sim_0,2), round(correct_0,2), round(correct_sim_0,2)))
# # ax = fig.add_subplot(122)
# # plt.hist(df.loc[(df['rt']!=-1)&(df['reward']==1),'rt'], bins=28, density=density, alpha=0.5)
# # plt.hist(df_sim_1['rt'], bins=28, density=density, alpha=0.5)
# # plt.title('yes: {} ({})\ncorrect: {} ({})'.format(round(yes_1,2), round(yes_sim_1,2), round(correct_1,2), round(correct_sim_1,2)))
# # fig.savefig(os.path.join('figs', 'bars_1.pdf'))



# #         # shell()

# #         # df['bins'] = pd.cut(df['time_target'], bins=[-1, 0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4, 4.75, 5.5, 6.5, 7.5, 9, 11, 15], labels=False)
# #         # df['start'] = df['bins'].map(dict(df.groupby('bins').mean()['time_target']))

# #         # shell()

# #         # for s, d in df.groupby('start'):
# #         #     plt.figure()
# #         #     plt.hist(d.loc[d['hit']==1,'rt'], bins=25)
# #         #     plt.hist(d.loc[d['fa']==1,'rt'], bins=25)

# #         # shell()

# #         if dataset == 'att_eff':
# #             if n_jobs == 1:
# #                 res = []
# #                 for params, d in df.groupby(['a', 'v', 'dc', 'k']):
# #                     res.append(fit_model(df=d.loc[:,['start', 'duration', 'response', 'rt']], att_eff=True, fit_bound=fit_bound, fit_leak=fit_leak))
# #             else:
# #                 res = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(fit_model)(df=d.loc[:,['start', 'duration', 'response', 'rt']], att_eff=True, fit_bound=fit_bound, fit_leak=fit_leak)
# #                             for params, d in df.groupby(['a', 'v', 'dc', 'k']))

# #             true = df.groupby(['a', 'v', 'dc', 'k']).mean().reset_index()[['a', 'v', 'dc', 'k']] / 100
# #             true.to_csv('data/true_{}_{}_{}.csv'.format(dataset, int(fit_bound),int(fit_leak)))
# #         elif dataset == 'empirical':
# #             res = []


# #             model_settings = {'fit': 'pyddm', 'depends_on': {'a':['reward'], 'u':None, 'v':['reward'], 't':['reward'], 'z':None, 'b':['reward'], 'k':None}, 
# #             'start_bias': False,  'drift_bias': True, 'leak': False, 'urgency': False, 'T_dur': 14}

# #             print('here we go!')
        
# #             params = pyddm_tools.fit_model_one_accumulator(df=df.loc[:,['start', 'duration', 'reward', 'response', 'rt']], model_settings=model_settings, subj_idx=0)
        
# #             res = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(fit_model_one_accumulator)(df=data.loc[:,['start', 'duration', 'response', 'rt']], 
# #                                                                                                         model_settings=model_settings, subj_idx=subj)
# #                         for subj, data in df.loc[(df['trial']%60)>=10].groupby(['subject_id']))

# #             shell()


# #             res.append(pyddm_tools.fit_model_one_accumulator(df=df.loc[:,['start', 'duration', 'reward', 'response', 'rt']], model_settings=model_settings, subj_idx=0))
    
# #         fitted = pd.concat(res)
# #         fitted.to_csv('data/fitted_{}_{}_{}.csv'.format(dataset, int(fit_bound), int(fit_leak)))
# #     else:
# #         df = pd.read_hdf('data/gonogo_1.h5', key='sims')
# #         df['stimulus'] = df['stimulus'].astype(int)
# #         df['response'] = df['response'].astype(int)
# #         df.loc[df['rt'].isna(), 'rt'] = -1
# #         df = df.loc[(df['k']==0),:]
# #         df = df.loc[df['a'].isin([ 50,   70,  90,  110, 130, 150]),:]
# #         df = df.loc[df['v'].isin([ 0,    50,  100, 150, 200, 250]),:]
# #         df = df.loc[df['dc'].isin([-100, -60, -20, 20,  60,  100]),:]

# #         res = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(fit_model)(df=d.loc[:,['stimulus', 'response', 'rt']], att_eff=False)
# #                     for params, d in df.groupby(['a', 'v', 'dc', 'k']))
# #         fitted = pd.concat(res)
# #         true = df.groupby(['a', 'v', 'dc', 'k']).mean().reset_index()[['a', 'v', 'dc', 'k']] / 100
    
# #         fitted.to_csv('data/fitted_gonogo.csv')
# #         true.to_csv('data/true_gonogo.csv')

# #     # df.loc[(df['a']==100)&(df['v']==100)&(df['dc']==0)&(df['k']==0),:].to_csv('data/gonogo_1.csv')
# #     # df = pd.read_csv('data/gonogo_1.csv')
# #     # df.loc[df['rt'].isna(), 'rt'] = -1

# #     # fit:
# #     # model_fit = fit_model(df=df.loc[:,['stimulus', 'response', 'rt']])

# # else:

# #     fit_bound = 0
# #     fit_leak = 1
# #     dataset = 'att_eff'
# #     fitted = pd.read_csv('data/fitted_{}_{}_{}.csv'.format(dataset, int(fit_bound),int(fit_leak)))
# #     true = pd.read_csv('data/true_{}_{}_{}.csv'.format(dataset, int(fit_bound),int(fit_leak)))

# #     # fitted['k'] = fitted['k']*fitted['v']

# #     colors = np.arange(fitted.shape[0])

# #     fig = plt.figure(figsize=(6,2))

# #     ax = fig.add_subplot(141)
# #     plt.plot(np.linspace(true['a'].min(), true['a'].max(), 100),
# #             np.linspace(true['a'].min(), true['a'].max(), 100),
# #             color='k', ls='--')
# #     if fit_bound:
# #         ax.scatter(true['a'], fitted['a'], c=colors)
# #     else:
# #         ax.scatter(true['a'], np.ones(len(true['a'])), c=colors)
# #     plt.xlabel('a (true)')
# #     plt.ylabel('a (fitted)')
# #     plt.title('Bound')

# #     ax = fig.add_subplot(142)
# #     plt.plot(np.linspace(true['v'].min(), true['v'].max(), 100),
# #             np.linspace(true['v'].min(), true['v'].max(), 100),
# #             color='k', ls='--')
# #     ax.scatter(true['v'], fitted['v'], c=colors)
# #     plt.xlabel('v (true)')
# #     plt.ylabel('v (fitted)')
# #     plt.title('Drift rate')

# #     ax = fig.add_subplot(143)
# #     plt.plot(np.linspace(true['dc'].min(), true['dc'].max(), 100),
# #             np.linspace(true['dc'].min(), true['dc'].max(), 100),
# #             color='k', ls='--')
# #     ax.scatter(true['dc'], fitted['b'], c=colors)
# #     plt.xlabel('v_bias (true)')
# #     plt.ylabel('v_bias (fitted)')
# #     plt.title('Drift bias')

# #     ax = fig.add_subplot(144)
# #     plt.axhline(true['k'].mean(), color='k', ls='--')
# #     if fit_leak:
# #         ax.scatter(true['k'], fitted['k'], c=colors)
# #     else:
# #         ax.scatter(true['k'], np.zeros(len(true['a'])), c=colors)
# #     plt.xlabel('k (true)')
# #     plt.ylabel('k (fitted)')
# #     plt.title('Leak')

# #     plt.tight_layout()
# #     fig.savefig('figs/fits_{}_{}_{}.pdf'.format(dataset, int(fit_bound),int(fit_leak)))

# #     shell()

# #     fit_bound = 0
# #     fit_leak = 1
# #     dataset = 'empirical'
# #     fitted = pd.read_csv('data/fitted_{}_{}_{}.csv'.format(dataset, int(fit_bound),int(fit_leak)))
# #     fitted['k'] = fitted['k']*fitted['v']
# #     fitted = fitted.drop(['Unnamed: 0', 'bic'], axis=1)

# #     fig = plt.figure(figsize=(2,2))
# #     ax = fig.add_subplot(111)
# #     sns.barplot(x='index', y=0, data=fitted.T.reset_index(), ax=ax)
# #     plt.ylim(-0.5,1.5)
# #     plt.tight_layout()
# #     fig.savefig('figs/fits_{}_{}_{}.pdf'.format(dataset, int(fit_bound),int(fit_leak)))

# #     shell()





# #     def plot_effect_parameters(df, param):

# #         colors = sns.color_palette("viridis", 11)

# #         fig = plt.figure(figsize=(6,3))
# #         ax = fig.add_subplot(121)
# #         for i, (p, data) in enumerate(df.groupby(param)):
# #             plt.hist(data.loc[data['hit']==1,'rt'], bins=100, histtype='step', cumulative=True, color=colors[i])
# #         plt.title(param)
# #         plt.xlabel('RT (s)')
# #         plt.ylabel('Trials')
# #         ax = fig.add_subplot(122)
# #         for i, (p, data) in enumerate(df.groupby(param)):
# #             plt.hist(data.loc[data['fa']==1,'rt'], bins=100, histtype='step', cumulative=True, color=colors[i])
# #         plt.title(str(df[param].unique()/100))
# #         plt.xlabel('RT (s)')
# #         plt.ylabel('Trials')
# #         plt.tight_layout()
    
# #         return fig

# #     for p in ['a', 'v', 'k', 'dc']:

# #         fig = plot_effect_parameters(df, param=p)
# #         fig.savefig('figs/effect_of_{}.pdf'.format(p))




# #     fig = plt.figure(figsize=(6,3))
# #     plt_nr = 0
# #     for a, d in df.groupby('v'):
# #         ax = fig.add_subplot(121)
# #         plt.hist(d.loc[d['hit']==1,'rt'], bins=100, histtype='step', cumulative=True, color=colors[plt_nr])
# #         ax = fig.add_subplot(122)
# #         plt.hist(d.loc[d['fa']==1,'rt'], bins=100, histtype='step', cumulative=True, color=colors[plt_nr])
# #         plt_nr += 1

# #     fig = plt.figure(figsize=(6,3))
# #     plt_nr = 0
# #     for a, d in df.groupby('dc'):
# #         ax = fig.add_subplot(121)
# #         plt.hist(d.loc[d['hit']==1,'rt'], bins=100, histtype='step', cumulative=True, color=colors[plt_nr])
# #         ax = fig.add_subplot(122)
# #         plt.hist(d.loc[d['fa']==1,'rt'], bins=100, histtype='step', cumulative=True, color=colors[plt_nr])
# #         plt_nr += 1

# #     fig = plt.figure(figsize=(6,3))
# #     plt_nr = 0
# #     for a, d in df.groupby('k'):
# #         ax = fig.add_subplot(121)
# #         plt.hist(d.loc[d['hit']==1,'rt'], bins=100, histtype='step', cumulative=True, color=colors[plt_nr])
# #         ax = fig.add_subplot(122)
# #         plt.hist(d.loc[d['fa']==1,'rt'], bins=100, histtype='step', cumulative=True, color=colors[plt_nr])
# #         plt_nr += 1