# coding: utf-8
import os
from functools import reduce
from pprint import pformat
import numpy as np
import scipy as sp
from scipy import stats
import pandas as pd
import statsmodels.api as sm
from statsmodels.stats.anova import AnovaRM
import json
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.transforms as transforms
import seaborn as sns
from joblib import Parallel, delayed
from tqdm import tqdm

from IPython import embed as shell

import analysis_tools

matplotlib.rcParams['pdf.fonttype'] = 42
matplotlib.rcParams['ps.fonttype'] = 42
sns.set(style='ticks', font='Arial', font_scale=1, rc={
    'axes.linewidth': 0.25, 
    'axes.labelsize': 7, 
    'axes.titlesize': 7, 
    'xtick.labelsize': 6, 
    'ytick.labelsize': 6, 
    'legend.fontsize': 6, 
    'xtick.major.width': 0.25, 
    'ytick.major.width': 0.25,
    'text.color': 'Black',
    'axes.labelcolor':'Black',
    'xtick.color':'Black',
    'ytick.color':'Black',} )
sns.plotting_context()
sns.set_palette("tab10")

plot_across_trials_full_session = 0
plot_simulations = 0
plot_survival = 0
plot_across_trials_from_block_type = 0
plot_across_blocks = 0
plot_individual_differences = 0
plot_pupil_behavior_correlations = 0
plot_cox_regression = 0
plot_mediation = 0
plot_state_occupancy_analysis = 0

def compute_local_utility(df, df_res, name, columns):
    
    columns_z = ['{}_z'.format(c) for c in columns]
    columns_d = ['signal:{}'.format(c) for c in columns]
    subj = df['subject_id'].iloc[0]
    
    for c in columns:
        df['{}_z'.format(c)] = (df[c]-df[c].mean()) / df[c].std()
    
    coefs_c = np.log(np.atleast_2d(df_res.loc[df_res['subject_id']==subj, columns].mean(axis=0)))
    coefs_d = np.log(np.atleast_2d(df_res.loc[df_res['subject_id']==subj, columns_d].mean(axis=0)))
    coefs = np.hstack((coefs_c, coefs_d))
    print(coefs)

    df['{}_c'.format(name)] = np.array((df.loc[:,columns_z] * coefs_c).sum(axis=1))
    df['{}_d'.format(name)] = np.array((df.loc[:,columns_z] * coefs_d).sum(axis=1))
    df['{}'.format(name)] = np.array((df.loc[:,columns_z+columns_z] * coefs).sum(axis=1))

    return df

def plot_delta_trial_density(df, pupil_measure, ax=None):

    kde_high = []
    kde_low = []
    walk_high = []
    walk_low = []
    counts_high = []
    counts_low = []

    x = np.linspace(0.1,0.9,100)
    x_binned_low = df.loc[df['reward']==0,:].groupby(['subject_id', '{}_bin'.format(pupil_measure)])[pupil_measure].mean().groupby(['{}_bin'.format(pupil_measure)]).mean().values.ravel()
    x_binned_high = df.loc[df['reward']==1,:].groupby(['subject_id', '{}_bin'.format(pupil_measure)])[pupil_measure].mean().groupby(['{}_bin'.format(pupil_measure)]).mean().values.ravel()
    x_binned = (x_binned_high + x_binned_low)/2
    for (s, r), d in df.groupby(['subject_id', 'reward']):

        kernel = sp.stats.gaussian_kde(d.loc[d['walk']==0, pupil_measure])
        if r == 0:
            kde_low.append(kernel(x))
            walk_low.append(d['walk'].sum() / d.shape[0])
        if r == 1:
            kde_high.append(kernel(x))
            walk_high.append(d['walk'].sum() / d.shape[0])

        counts = d.groupby(['{}_bin'.format(pupil_measure)])['trial'].count().reset_index()
        counts['trial'] = counts['trial'] / sum(counts['trial']) * 100
        counts['subject_id'] = s
        if r == 0:
            counts_low.append(counts)
        if r == 1:
            counts_high.append(counts)

    counts_high = pd.concat(counts_high)
    counts_high = counts_high.set_index(['subject_id', '{}_bin'.format(pupil_measure)])
    counts_low = pd.concat(counts_low)
    counts_low = counts_low.set_index(['subject_id', '{}_bin'.format(pupil_measure)])
    counts = (counts_high + counts_low) / 2
    counts_diff = (counts_high - counts_low)

    counts_high = counts_high.reset_index()
    counts_low = counts_low.reset_index()

    kde_low = np.vstack(kde_low)
    kde_high = np.vstack(kde_high)


    fig = plt.figure(figsize=(3.5,1.75))
    ax = fig.add_subplot(121)

    mean = kde_low.mean(axis=0)
    sem = sp.stats.sem(kde_low, axis=0)
    plt.fill_between(x, mean-sem, mean+sem, alpha=0.2)
    plt.plot(x,mean)

    mean = kde_high.mean(axis=0)
    sem = sp.stats.sem(kde_high, axis=0)
    plt.fill_between(x, mean-sem, mean+sem, alpha=0.2)
    plt.plot(x,mean)
    plt.xlabel('Pupil size (% max)')
    plt.ylabel('Prob. density')

    ax = ax.twinx()
    plt.errorbar(x_binned_low[-1], 
                    counts_low.loc[counts_low['{}_bin'.format(pupil_measure)]==100,'trial'].mean(), 
                    yerr=counts_low.loc[counts_low['{}_bin'.format(pupil_measure)]==100,'trial'].sem(), 
                    fmt='x', color=sns.color_palette()[0])
    plt.errorbar(x_binned_high[-1], 
                    counts_high.loc[counts_high['{}_bin'.format(pupil_measure)]==100,'trial'].mean(), 
                    yerr=counts_high.loc[counts_high['{}_bin'.format(pupil_measure)]==100,'trial'].sem(), 
                    fmt='x', color=sns.color_palette()[1])

    
    plt.ylabel('Trial occupancy (% of trials)')

    ax = fig.add_subplot(122)

    mean = counts_diff.groupby(['{}_bin'.format(pupil_measure)]).mean().values.ravel()
    sem = counts_diff.groupby(['{}_bin'.format(pupil_measure)]).sem().values.ravel()

    plt.errorbar(x_binned[:-1], mean[:-1], yerr=sem[:-1], fmt='-o', color='green')
    plt.errorbar(x_binned[-1], mean[-1], yerr=sem[-1], fmt='x', color='green') #, color=sns.color_palette()[2])


    plt.axhline(0, color='k', lw=0.5)
    plt.xlabel('Pupil size (% max)')
    plt.ylabel('Δ Trial occupancy (high vs low)')

    sns.despine(trim=False)
    plt.tight_layout()

    return fig

def compute_trial_in_block(df):
    df['trial_in_block'] = np.arange(df.shape[0]) + (60*df['block_id'].iloc[0])
    return df

####################################

# project info:
project_dir = '/Users/jwdegee/Library/CloudStorage/Box-Box/McGinleyLabFiles/projects/2021_Attentional_Effort'
fig_dir = os.path.join(project_dir, 'figs')
data_dir = os.path.join(project_dir, 'data')
rt_cutoff = 0.15
n_jobs = 8
# n_boot = 25
n_boot = 500
exp = 'behavior'
# exp = 'psychometric'

# load:
df = pd.read_csv(os.path.join(data_dir, 'data_{}.csv'.format(exp)))
df = df.groupby(['subject_id', 'session_id', 'block_id'], group_keys=False).apply(compute_trial_in_block)
df = df.drop(columns=['Unnamed: 0'])
df = df.loc[~(df['noise_dur']==np.inf),:]
df = df.drop(['date', 'start_time'], axis=1)
df['hit'] = (df['outcome']==0).astype(int)
df['miss'] = (df['outcome']==1).astype(int)
df['fa'] = (df['outcome']==2).astype(int)
df['cr'] = (df['outcome']==3).astype(int)
df['correct'] = ((df['hit']==1)|(df['cr']==1)).astype(int)
df['choice'] = ((df['hit']==1)|(df['fa']==1)).astype(int)
df.loc[df['outcome']==0, 'rt2'] = df.loc[df['outcome']==0, 'rt']-df.loc[df['outcome']==0, 'noise_dur']
df['reward_ind'] = (df['trial_since_hit']>=15) 
df.loc[df['outcome']==3, 'noise_dur'] = 14
df['catch_session'] = df.groupby(['subject_id', 'session_id'])['catch'].transform(np.max)
df['reward_obtained'] = 0
df.loc[(df['correct']==1)&(df['reward']==0), 'reward_obtained'] = 2
df.loc[(df['correct']==1)&(df['reward']==1), 'reward_obtained'] = 12
df['reward_obtained2'] = 0
df.loc[(df['correct']==1)&(df['reward']==0), 'reward_obtained2'] = 1
df.loc[(df['correct']==1)&(df['reward']==1), 'reward_obtained2'] = 2

# add previous trial info:
df = df.groupby(['subject_id', 'session_id']).apply(analysis_tools.add_previous_trial_info, 
                                                  columns=['hit', 'fa', 'choice', 'correct', 'walk', 
                                                           'reward_obtained', 'reward_obtained2']).reset_index(drop=True)
df = df.loc[df['previous_trial_exists'],:]

# # add next trial info:
# df = df.groupby(['subject_id', 'block_id']).apply(analysis_tools.add_next_trial_info, columns=['hit', 'choice', 'correct']).reset_index(drop=True)
# df = df.loc[df['next_trial_exists'],:]

df['too_fast'] = df['rt']<rt_cutoff
print(df.loc[df['reward_ind']&df['block_id']>0,:].groupby(['subject_id', 'reward'])['too_fast'].mean().groupby(['reward']).mean())
# sns.pointplot(data=df.loc[df['reward_ind']&df['block_id']>0,:], x='reward', y='too_fast', units='subject_id', errorbar='se')

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

# correct baseline pupil size:
def correct_baseline_pupil_responses(df):
    import statsmodels.formula.api as smf
    
    ind = (df['walk']==0) & (df['block_id']>0) & (df['reward_ind'])
    ind2 = (df['walk']==1) & (df['block_id']>0) & (df['reward_ind'])
    
    # model 1:
    model = smf.ols(formula='pupil_trial_start ~ block_id', data=df.loc[ind,:]).fit()
    df['pupil_trial_start_c1'] = np.NaN
    df.loc[ind, 'pupil_trial_start_c1'] = model.resid + df.loc[ind, 'pupil_trial_start'].mean()

    # model 2:
    model = smf.ols(formula='pupil_trial_start ~ hit_p', data=df.loc[ind,:]).fit()
    df['pupil_trial_start_c2'] = np.NaN
    df.loc[ind, 'pupil_trial_start_c2'] = model.resid + df.loc[ind, 'pupil_trial_start'].mean()

    # model 3:
    model = smf.ols(formula='pupil_trial_start ~ hit_p + block_id', data=df.loc[ind,:]).fit()
    df['pupil_trial_start_c3'] = np.NaN
    df.loc[ind, 'pupil_trial_start_c3'] = model.resid + df.loc[ind, 'pupil_trial_start'].mean()
    
    # add back walking:
    df.loc[ind2, 'pupil_trial_start_c1'] = df.loc[ind2, 'pupil_trial_start']
    df.loc[ind2, 'pupil_trial_start_c2'] = df.loc[ind2, 'pupil_trial_start']
    df.loc[ind2, 'pupil_trial_start_c3'] = df.loc[ind2, 'pupil_trial_start']

    return df

df = df.groupby(['subject_id', 'reward'], group_keys=False).apply(correct_baseline_pupil_responses)
df['pupil_trial_start_c3nowalk'] = df['pupil_trial_start_c3'].copy()

# add pupil bins:
# bins = [-10, 0.15, 0.25, 0.3, 0.45, 0.5, 0.6, 0.7, 10]
bins = [-10, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.6, 0.7, 0.8, 10]
ind = (df['block_id']>0)&(df['reward_ind'])
df.loc[ind&(df['walk']==0), 'pupil_trial_start_bin'] = df.loc[ind&(df['walk']==0),:].groupby(['subject_id'], group_keys=False)['pupil_trial_start'].apply(pd.cut, bins=bins, labels=False).astype(int)
df.loc[ind&(df['walk']==1), 'pupil_trial_start_bin'] = 100
df.loc[ind&(df['walk']==0), 'pupil_trial_start_c1_bin'] = df.loc[ind&(df['walk']==0),:].groupby(['subject_id'], group_keys=False)['pupil_trial_start_c1'].apply(pd.cut, bins=bins, labels=False).astype(int)
df.loc[ind&(df['walk']==1), 'pupil_trial_start_c1_bin'] = 100
df.loc[ind&(df['walk']==0), 'pupil_trial_start_c2_bin'] = df.loc[ind&(df['walk']==0),:].groupby(['subject_id'], group_keys=False)['pupil_trial_start_c2'].apply(pd.cut, bins=bins, labels=False).astype(int)
df.loc[ind&(df['walk']==1), 'pupil_trial_start_c2_bin'] = 100
df.loc[ind&(df['walk']==0), 'pupil_trial_start_c3_bin'] = df.loc[ind&(df['walk']==0),:].groupby(['subject_id'], group_keys=False)['pupil_trial_start_c3'].apply(pd.cut, bins=bins, labels=False).astype(int)
df.loc[ind&(df['walk']==1), 'pupil_trial_start_c3_bin'] = 100
df.loc[ind,                 'pupil_trial_start_c3nowalk_bin'] = df.loc[ind,:].groupby(['subject_id'], group_keys=False)['pupil_trial_start_c3nowalk'].apply(pd.cut, bins=bins, labels=False).astype(int)

df['state'] = 0
df.loc[(df['pupil_trial_start_bin']>=3)&(df['pupil_trial_start_bin']<=10), 'state'] = 1
df.loc[(df['pupil_trial_start_bin']>=11)&(df['pupil_trial_start_bin']<=12), 'state'] = 2
df.loc[(df['pupil_trial_start_bin']==100), 'state'] = 3

# save for ddm:
print('saving for DDM...')
df.to_csv(os.path.join(data_dir, 'data_ddm_{}.csv'.format(exp)))

# fast hits to false alarms:
df = analysis_tools.fast_hits_to_false_alarms(df, rt_cutoff=rt_cutoff)

# add distance from optimal:
ind = (df['block_id']>0)&(df['reward_ind'])
optimal_arousal = float(df.loc[ind&(df['pupil_trial_start_c3_bin']==4),:].groupby(['subject_id', 'reward'])['pupil_trial_start_c3'].mean().groupby(['subject_id']).mean().mean())
print(optimal_arousal)
df['dist_from_optimal'] = (df['pupil_trial_start_c3'] - optimal_arousal).abs()

# assert(df['noise_dur'].max() == 11)
assert(df['rt'].min() >= 0)

# print:
nr_subjects = df.groupby(['subject_id']).count().shape[0]
nr_sessions = df.groupby(['subject_id', 'session_id']).count().shape[0]
print('nr subjects = {}'.format(nr_subjects))
print('nr sessions = {}'.format(nr_sessions))
print(df.groupby(['subject_id', 'session_id'])['trial'].count().groupby(['subject_id']).count().min())
print(df.groupby(['subject_id', 'session_id'])['trial'].count().groupby(['subject_id']).count().max())
print(df.groupby(['subject_id'])['trial'].count().min())
print(df.groupby(['subject_id'])['trial'].count().max())
print(df.shape)

# print(df.groupby(['subject_id', 'session_id'])['time_trial_end'].max().sum()/60/60)
# plt.figure()
# plt.hist(df.loc[df['walk']==0,'pupil_trial_start'], bins=100)
# plt.hist(df.loc[df['walk']==1,'pupil_trial_start'], bins=100)
# for b in bins[1:-1]:
#     plt.axvline(b, color='r')

#################################
# plot across the full session:
#################################
if plot_across_trials_full_session:

    n_boot = 350
    os.environ['JOBLIB_TEMP_FOLDER'] = '/tmp'
    for plot_type in [0]:

        if plot_type == 0:
            xs = ['trial_in_block', 'trial_since_hit_s']
            res1 = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(analysis_tools.compute_metrics)
                                                                                (df=df.groupby(['subject_id', 'block_id']).sample(frac=1, replace=True), 
                                                                                groupby=['block_id', 'trial_in_block']) for i in tqdm(range(n_boot)))
            res2 = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(analysis_tools.compute_metrics)
                                                                                (df=df.loc[(df['trial_since_hit']>=2) & (df['trial_since_hit']<50),:].groupby(['subject_id', 'block_id']).sample(frac=1, replace=True), 
                                                                                groupby=['block_id', 'trial_since_hit_s']) for i in tqdm(range(n_boot)))

        # res1 = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(analysis_tools.compute_metrics)
        #                                                                     (df=df.loc[(df['trial_since_hit']>=2) & (df['reward_obtained2_p2']!=np.NaN) & (df['reward_obtained2_p2']==0) & (df['trial_since_hit']<50),:].groupby(['subject_id', 'block_id']).sample(frac=1, replace=True), 
        #                                                                     groupby=['block_id', 'trial_since_hit_s']) for i in tqdm(range(n_boot)))

        # res2 = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(analysis_tools.compute_metrics)
        #                                                                     (df=df.loc[(df['trial_since_hit']>=2) & (df['reward_obtained2_p2']!=np.NaN) & (df['reward_obtained2_p2']!=0) & (df['trial_since_hit']<50),:].groupby(['subject_id', 'block_id']).sample(frac=1, replace=True), 
        #                                                                     groupby=['block_id', 'trial_since_hit_s']) for i in tqdm(range(n_boot)))

        elif plot_type == 1:
            xs = ['trial_since_hit_s', 'trial_since_hit_s']
            res1 = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(analysis_tools.compute_metrics)
                                                                                (df=df.loc[(df['trial_since_hit']>=1) & (df['reward_obtained2_p1']!=np.NaN) & (df['reward_obtained2_p1']==0) & (df['trial_since_hit']<50),:].groupby(['subject_id', 'block_id']).sample(frac=1, replace=True), 
                                                                                groupby=['block_id', 'trial_since_hit_s']) for i in tqdm(range(n_boot)))

            res2 = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(analysis_tools.compute_metrics)
                                                                                (df=df.loc[(df['trial_since_hit']>=1) & (df['reward_obtained2_p1']!=np.NaN) & (df['reward_obtained2_p1']!=0) & (df['trial_since_hit']<50),:].groupby(['subject_id', 'block_id']).sample(frac=1, replace=True), 
                                                                                groupby=['block_id', 'trial_since_hit_s']) for i in tqdm(range(n_boot)))

        dfs1 = []
        for i,r in enumerate(res1):
            r['subject_id'] = i
            dfs1.append(r)
        df_res1 = pd.concat(dfs1)
        dfs2 = []
        for i,r in enumerate(res2):
            r['subject_id'] = i
            dfs2.append(r)
        df_res2 = pd.concat(dfs2)

        # offsets = df.loc[(df['trial_since_hit']==0), :].groupby(['subject_id', 'block_id'])['trial'].mean().groupby(['block_id']).mean().round()
        # for t, o in zip([0,60,120,180,240,300,360,420], offsets):
        #     ind = (df_res2['trial_since_hit_s']>=t)&(df_res2['trial_since_hit_s']<(t+60))
        #     df_res2.loc[ind, 'trial_since_hit_s'] = df_res2.loc[ind, 'trial_since_hit_s'] + (o-df_res2.loc[ind, 'trial_since_hit_s'].min())

        for m in ['hit', 'rt2', 'tr_hr', 'tr_far', 'tr_d', 'tr_c',
                'pupil_trial_start', 'pupil_trial_start_c3', 'walk',
                'dist_from_optimal']:
            fig = analysis_tools.plot_baselines_across_trials(dfs=[df_res1, df_res2], 
                                                            xs=xs, 
                                                            measure=m, error='ci', scale_sem=1)
            fig.savefig(os.path.join(fig_dir, '{}_{}_across_{}_{}.pdf'.format(exp, m, 'session', plot_type)))


        for m in ['pupil_trial_start_c3', 'dist_from_optimal']:
            fig = analysis_tools.plot_baselines_across_trials(dfs=[df_res1.loc[df_res1['trial_since_hit']>=15,:], 
                                                                   df_res2.loc[df_res2['trial_since_hit']>=15,:]], 
                                                            xs=xs, 
                                                            measure=m, error='ci', scale_sem=1)
            fig.savefig(os.path.join(fig_dir, '{}_{}_across_{}_{}.pdf'.format(exp, m, 'session', plot_type)))



    # # pupil lick:
    # for m in ['pupil_lick', 'pupil_lick_c', 'pupil_lick_s', 'pupil_lick_s_c']:
    #     for trial_type in ['choice', 'hit', 'fa']:
    #         df_res_1 = df.loc[df[trial_type]==1,:].groupby(['subject_id', 'block_id', 'trial_s', 'outcome']).mean().reset_index().groupby(['subject_id', 'trial_s']).mean().reset_index()
    #         df_res_2 = df.loc[df[trial_type]==1,:].loc[(df['trial_since_hit']>=2) & (df['trial_since_hit']<50), :].groupby(
    #                                                                                         ['subject_id', 'block_id', 'trial_since_hit_s', 'outcome']).mean().reset_index().groupby(
    #                                                                                         ['subject_id', 'block_id', 'trial_since_hit_s']).mean().reset_index()
    #         offsets = df.loc[(df['trial_since_hit']==0), :].groupby(['subject_id', 'block_id'])['trial'].mean().groupby(['block_id']).mean().round()
    #         for t, o in zip([60,120,180,240,300,360,420], offsets):
    #             ind = (df_res_2['trial_since_hit_s']>=t)&(df_res_2['trial_since_hit_s']<(t+60))
    #             df_res_2.loc[ind, 'trial_since_hit_s'] = df_res_2.loc[ind, 'trial_since_hit_s'] + (o-df_res_2.loc[ind, 'trial_since_hit_s'].min())
    #         fig = analysis_tools.plot_baselines_across_trials(dfs=[df_res_1, df_res_2], xs=['trial_s', 'trial_since_hit_s'], measure=m, scale_sem=1)
    #         fig.savefig(os.path.join(fig_dir, '{}_{}_across_{}_{}.pdf'.format(exp, m, 'session', trial_type)))

#################################
# exclude first block:
#################################
df = df.loc[df['block_id']>0,:].reset_index(drop=True)

# reward rate:
res = analysis_tools.compute_metrics(df.loc[df['reward_ind']], groupby=['subject_id', 'block_id'])
print(res.groupby(['subject_id', 'reward']).mean().groupby(['reward']).mean()['hit'])
print(res.groupby(['subject_id', 'reward']).mean().groupby(['reward']).sem()['hit'])

print(res.loc[res['block_id']==2,'hit'].mean())
print(res.loc[res['block_id']==2,'hit'].sem())
print(res.loc[res['block_id']==6,'hit'].mean())
print(res.loc[res['block_id']==6,'hit'].sem())

#################################
# plot simulations:
#################################
if plot_simulations:

    dfs = []
    for rate in np.linspace(0.05,1,40):
        df_sim = analysis_tools.simulate_data(rate=rate, nr_trials=500000)
        df_sim['rate'] = rate
        dfs.append(df_sim)
    df_sim = pd.concat(dfs)
    df_sim['reward'] = 1
    df_sim['catch'] = 0

    df_res = analysis_tools.compute_metrics(df_sim, groupby=['rate'])

    optimal = float(df_res.loc[df_res['hit']==df_res['hit'].max(), 'rate'].iloc[0])

    fig = plt.figure(figsize=(5,1.5))
    ax = fig.add_subplot(141)
    plt.plot(df_res['rate'], df_res['tr_c'])
    plt.axvline(optimal, color='green')
    ax.set_xscale('log')
    plt.xlabel('Rate')
    plt.ylabel('Bias (s.d.)')
    lim = ax.get_ylim()
    ax = fig.add_subplot(142)
    plt.plot(df_res['rate'], df_res['tr_d'])
    plt.axvline(optimal, color='green')
    ax.set_xscale('log')
    plt.xlabel('Rate')
    plt.ylabel('Sensitivity (s.d.)')
    ax.set_ylim(lim)
    ax = fig.add_subplot(143)
    plt.plot(df_res['rate'], df_res['rt2'])
    plt.axvline(optimal, color='green')
    ax.set_xscale('log')
    plt.xlabel('Rate')
    plt.ylabel('RT (s)')
    ax = fig.add_subplot(144)
    plt.plot(df_res['rate'], df_res['hit'])
    plt.axvline(optimal, color='green')
    ax.set_xscale('log')
    

    res = df.loc[df['reward_ind'],:].groupby(['subject_id', 'reward']).apply(analysis_tools.sdt_ps).reset_index()

    # plt.axvline(res.loc[res['reward']==0, 'r_ps'].mean(), color=sns.color_palette()[0])
    # plt.axvline(res.loc[res['reward']==1, 'r_ps'].mean(), color=sns.color_palette()[1])
    plt.xlabel('Rate')
    plt.ylabel('P(reward)')
    sns.despine(trim=False)
    plt.tight_layout()
    fig.savefig(os.path.join(fig_dir, '{}_simulations.pdf'.format(exp)))

#################################
# plot survival:
#################################
if plot_survival:
    fig = analysis_tools.kaplan_meier_plots(df.loc[df['reward_ind'],:].copy(), 
                                            nbins=10, dt=0.1)
    fig.savefig(os.path.join(fig_dir, '{}_survival.pdf'.format(exp)))

#################################
# plot across trials from # hit:
#################################
if plot_across_trials_from_block_type:

    # for trial_measure in ['trial_since_hit_s', 'hit_block_s']:
    for trial_measure in ['trial_since_hit']:

        groupby = ['reward', trial_measure]
        trial_max = 50

        # res = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(analysis_tools.compute_metrics)
        #                                                                     (df=df.loc[(df[trial_measure]>=2) & (df[trial_measure]<trial_max),:].groupby(['subject_id', 'reward']).sample(frac=1, replace=True), 
        #                                                                     groupby=groupby) for i in tqdm(range(n_boot)))


        res1 = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(analysis_tools.compute_metrics)
                                                                            (df=df.loc[(df['trial_since_hit']>=1) & (df['reward_obtained2_p1']!=np.NaN) & (df['reward_obtained2_p1']==0) & (df['trial_since_hit']<50),:].groupby(['subject_id', 'block_id']).sample(frac=1, replace=True), 
                                                                            groupby=groupby) for i in tqdm(range(n_boot)))

        res2 = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(analysis_tools.compute_metrics)
                                                                            (df=df.loc[(df['trial_since_hit']>=1) & (df['reward_obtained2_p1']!=np.NaN) & (df['reward_obtained2_p1']!=0) & (df['trial_since_hit']<50),:].groupby(['subject_id', 'block_id']).sample(frac=1, replace=True), 
                                                                            groupby=groupby) for i in tqdm(range(n_boot)))

        shell()

        dfs1 = []
        for i,r in enumerate(res1):
            r['subject_id'] = i
            dfs1.append(r)
        df_res1 = pd.concat(dfs1)
        dfs2 = []
        for i,r in enumerate(res2):
            r['subject_id'] = i
            dfs2.append(r)
        df_res2 = pd.concat(dfs2)

        for m in ['hit', 'rt2', 'tr_hr', 'tr_far', 'tr_d', 'tr_c',
                'pupil_trial_start', 'pupil_trial_start_c3', 'walk',
                'dist_from_optimal']:
            # fig, time_constants, percentages = analysis_tools.plot_behavior_since_hit(dfs=[df_res1.loc[~df_res1[m].isna(),:], df_res2.loc[~df_res2[m].isna(),:]], 
            #                                                                           x_measure=trial_measure, y_measure=m, line=15, error='ci', scale_sem=1)
            # fig.savefig(os.path.join(fig_dir, '{}_since_hit_{}_{}.pdf'.format(exp, trial_measure, m)))
            # print(percentages)

            fig = analysis_tools.plot_behavior_since_hit_mean(dfs=[df_res1.loc[~df_res1[m].isna(),:], df_res2.loc[~df_res2[m].isna(),:]], 
                                                                                      x_measure=trial_measure, y_measure=m, line=15, error='ci', scale_sem=1)
            fig.savefig(os.path.join(fig_dir, '{}_since_hit_{}_{}.pdf'.format(exp, trial_measure, m)))

        # for y_measure in ['r_ps', 'd_ps']:
        #     n_jobs = 64
        #     globals().update(locals())
        #     res = Parallel(n_jobs=n_jobs, verbose=1, backend='multiprocessing')(delayed(analysis_tools.plot_behavior_since_hit)
        #                                 (df=df_res.groupby(['reward', 'trial_since_hit']).sample(frac=1, replace=True), 
        #                                 x_measure=trial_measure, y_measure=y_measure, line=15, scale_sem=1, plot=False) for i in tqdm(range(1000)))
        #     globals().update(locals())
        #     percentages = np.vstack([res[i][2] for i in range(len(res))])
        #     time_constants = np.vstack([res[i][1] for i in range(len(res))])

        #     print()
        #     print(percentages[:,0:3].mean(axis=0))
        #     print(np.round((np.percentile(percentages[:,0:3], 84, axis=0) - np.percentile(percentages[:,0:3], 16, axis=0)) / 2, 2))
        #     print(np.round((np.percentile(time_constants[:,1], 84, axis=0) - np.percentile(time_constants[:,1], 16, axis=0)) / 2, 2))
        #     print()
        #     print(percentages[:,3:6].mean(axis=0))
        #     print(np.round((np.percentile(percentages[:,3:6], 84, axis=0) - np.percentile(percentages[:,3:6], 16, axis=0)) / 2, 2))
        #     print(np.round((np.percentile(time_constants[:,0], 84, axis=0) - np.percentile(time_constants[:,0], 16, axis=0)) / 2, 2))

#################################
# plot across blocks:
#################################
if plot_across_blocks:
    
    # main analysis:
    groupby = ['subject_id', 'block_id']
    df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'],:], groupby=groupby)
    df_res['reward'] = df_res['block_id']%2
    for m in ['hit', 'fa', 'miss', 'tr_hr', 'tr_far', 'tr_d', 'tr_c', 'rt2', 'reward_ml',
              'pupil_trial_start', 'pupil_trial_start_c1', 'pupil_trial_start_c2', 
              'pupil_trial_start_c3', 'dist_from_optimal', 'walk',]:
    # for m in ['tr_d_predicted']:
        fig = analysis_tools.plot_scalars_across_blocks(df=df_res, x='block_id', y=m)
        fig.savefig(os.path.join(fig_dir, '{}_sdt_across_blocks_{}.pdf'.format(exp,m)))

    # use all trials within a block:
    groupby = ['subject_id', 'block_id']
    df_res = analysis_tools.compute_metrics(df=df.loc[:,:], groupby=groupby)
    df_res['reward'] = df_res['block_id']%2
    for m in ['hit', 'tr_hr', 'tr_far', 'tr_d', 'tr_c', 'rt2', 'reward_ml']:
        fig = analysis_tools.plot_scalars_across_blocks(df=df_res, x='block_id', y=m)
        fig.savefig(os.path.join(fig_dir, '{}_sdt_across_blocks_alltrials_{}.pdf'.format(exp,m)))

    # control for previous hit:
    groupby = ['subject_id', 'block_id', 'hit_p']
    df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'],:], groupby=groupby)
    df_res['reward'] = df_res['block_id']%2
    for m in ['hit', 'tr_hr', 'tr_far', 'tr_d', 'tr_c', 'rt2', 'reward_ml',]:
    # for m in ['tr_c', 'tr_d',  'rt2', 'hit',]:
        fig = analysis_tools.plot_scalars_across_blocks_ptrial(df=df_res, x='block_id', y=m)
        fig.savefig(os.path.join(fig_dir, '{}_sdt_across_blocks_ptrial_{}.pdf'.format(exp,m)))

    # control for time on task:
    groupby = ['subject_id', 'block_id']
    df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'],:], groupby=groupby)
    df_res['reward'] = df_res['block_id']%2
    df_res_c = df_res.loc[(df_res['block_id']>1)&(df_res['block_id']<6),:]
    df_res_cc = df_res.loc[(df_res['block_id']>2)&(df_res['block_id']<5),:]
    df_res.loc[df_res['block_id']==1, 'block'] = 1
    df_res.loc[df_res['block_id']==2, 'block'] = 2
    df_res.loc[df_res['block_id']==3, 'block'] = 3
    df_res.loc[df_res['block_id']==4, 'block'] = 4
    df_res.loc[df_res['block_id']==5, 'block'] = 3
    df_res.loc[df_res['block_id']==6, 'block'] = 4
    df_res_c.loc[df_res_c['block_id']==2, 'block'] = 1
    df_res_c.loc[df_res_c['block_id']==3, 'block'] = 1
    df_res_c.loc[df_res_c['block_id']==4, 'block'] = 2
    df_res_c.loc[df_res_c['block_id']==5, 'block'] = 4
    df_res_cc.loc[df_res_cc['block_id']==3, 'block'] = 2
    df_res_cc.loc[df_res_cc['block_id']==4, 'block'] = 3
    df_res = pd.concat((df_res, df_res_c, df_res_cc))
    df_res = df_res.groupby(['subject_id', 'reward', 'block']).mean().reset_index()
    for m in ['hit', 'tr_hr', 'tr_far', 'tr_d', 'tr_c', 'rt2', 'reward_ml']:
        fig = analysis_tools.plot_scalars_across_blocks(df=df_res, x='block', y=m)
        fig.savefig(os.path.join(fig_dir, '{}_sdt_across_blocks_timetask_{}.pdf'.format(exp,m)))

    # catch session control:
    groupby = ['subject_id', 'block_id']
    for c in [0,1]:
        df1 = df.loc[(df['reward_ind']) & (df['catch_session']==c),:]
        trials = df1.groupby(['subject_id', 'session_id'])['trial'].count().reset_index()
        trials = trials.loc[trials['trial']>=(200-1),:]
        df1 = df1[df1.set_index(['subject_id','session_id']).index.isin(trials.set_index(['subject_id','session_id']).index)].reset_index(drop=True)
        groupby = ['subject_id', 'block_id']
        df_res = analysis_tools.compute_metrics(df=df1.loc[df1['reward_ind'],:], groupby=groupby)
        df_res['reward'] = df_res['block_id']%2
        for m in ['hit', 'tr_hr', 'tr_far', 'tr_d', 'tr_c', 'rt2', 'reward_ml']:
            fig = analysis_tools.plot_scalars_across_blocks(df=df_res, x='block_id', y=m)
            fig.savefig(os.path.join(fig_dir, '{}_sdt_across_blocks_catch_{}_{}.pdf'.format(exp,c,m)))

    # # across correlations:

    # groupby = ['subject_id', 'block_id']
    # df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'],:], 
    #                                             groupby=groupby)
    # df_res['reward'] = df_res['block_id']%2
    # df_res['dummy'] = 1

    # # # simulate data:
    # # df_sim = df.loc[df['reward_ind'],:].groupby(groupby).apply(analysis_tools.simulate_data, method='licks_ps').reset_index()
    # # df_sim_cox = behavior.cox_prepare_data(df_sim)

    # # # metrics simulated data:
    # # df_sim_sdt = df_sim_cox.groupby(groupby).apply(analysis_tools.sdt).reset_index()
    # # df_sim_sdt_ps = df_sim_cox.groupby(groupby).apply(analysis_tools.sdt_ps).reset_index()
    # # df_sim_res = df_sim.groupby(groupby).mean().reset_index()
    # # df_sim_res['hit'] = df_sim_res['hit'] * 100
    # # df_sim_res['rt'] = df_sim.groupby(groupby).median().reset_index()['rt']
    # # # df_sim_res = df_sim_sdt.merge(df_sim_res, on=groupby)
    # # globals().update(locals())
    # # df_sim_res = reduce(lambda left, right: pd.merge(left,right, on=groupby), [df_sim_sdt, df_sim_sdt_ps, df_sim_res])
    # # df_sim_res['reward'] = ((df_sim_res['block_id']==1)|(df_sim_res['block_id']==3)|(df_sim_res['block_id']==5)).astype(int)

    # # df_res_corrected = (df_res.set_index(groupby) - df_sim_res.set_index(groupby)).reset_index()
    # # df_res_corrected['reward'] = ((df_res_corrected['block_id']==1)|(df_res_corrected['block_id']==3)|(df_res_corrected['block_id']==5)).astype(int)

    # for y in ['fa', 'hit', 'rt', 'fa_rt', 'tr_hr', 'tr_far', 'tr_d', 'tr_c',]:
    #             # 'pupil_trial_start', 'pupil_trial_start_c1', 'pupil_trial_start_ct', 'pupil_trial_start_std', 
    #             # 'pupil_trial_start_c1_std', 'pupil_trial_start_c1_stability', 'walk', 'dist_from_optimal', 'dist_from_optimal_ct']:
    #     fig = analysis_tools.plot_scalars_across_blocks(df=df_res, x='block_id', y=y)
    #     fig.savefig(os.path.join(fig_dir, '{}_sdt_across_blocks_{}.pdf'.format(exp,y)))

    # for y in ['pupil_trial_start_ct', 'dist_from_optimal_ct']:
    #     plt.figure(figsize=(1.5,1.75))
    #     sns.pointplot(x='dummy', y=y, hue='reward', units='subject_id', ci=68, data=df_res)

    #     # if not (('pupil' in y) | ('walk' in y)):

    #     #     fig = analysis_tools.plot_scalars_across_blocks(df=df_sim_res, x='block_id', y=y)
    #     #     fig.savefig(os.path.join(fig_dir, '{}_sdt_across_blocks_{}_sim.pdf'.format(exp,y)))

    #     #     fig = analysis_tools.plot_scalars_across_blocks(df=df_res_corrected, x='block_id', y=y)
    #     #     fig.savefig(os.path.join(fig_dir, '{}_sdt_across_blocks_{}_corrected.pdf'.format(exp,y)))

    # for trial_type in ['choice', 'hit', 'fa']:
    #     df_res = df.loc[df['reward_ind'] & (df[trial_type]==1),:].groupby(['subject_id', 'block_id', 'outcome']).mean().reset_index().groupby(groupby).mean().reset_index()
    #     for y in ['pupil_lick', 'pupil_lick_c', 'pupil_lick_s', 'pupil_lick_s_c']:
    #         fig = analysis_tools.plot_scalars_across_blocks(df=df_res, x='block_id', y=y)
    #         fig.savefig(os.path.join(fig_dir, '{}_sdt_across_blocks_{}_{}.pdf'.format(exp, y, trial_type)))

    # df['coh_bin'] = 5 - df.groupby(['subject_id', 'session_id'])['difficulty'].apply(pd.cut, 6, labels=False).values
    # df['coh_bin2'] = np.floor(0.5*(df['coh_bin']))

    # if exp == 'psychometric':
    #     groupby = ['subject_id', 'session_id', 'reward', 'coh_bin2']

    #     df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'],:], groupby=groupby)
    #     for y in ['hit', 'tr_d']:
    #     # for y in ['rr']:
    #         fig = analysis_tools.plot_scalars_across_blocks(df=df_res, x='coh_bin2', y=y, collapse=False)
    #         fig.savefig(os.path.join(fig_dir, '{}_sdt_across_blocks_{}_reward.pdf'.format(exp, y)))

    #         fig = analysis_tools.plot_scalars_across_blocks(df=df_res, x='coh_bin2', y=y, collapse=True)
    #         fig.savefig(os.path.join(fig_dir, '{}_sdt_across_blocks_{}.pdf'.format(exp, y)))

    # # pull across hits / false alarms:
    # groupby = ['subject_id', 'block_id']
    # df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'] & (df['hit']==1),:], 
    #                                             df_cox=df_cox.loc[df_cox['reward_ind'] & (df_cox['hit']==1),:], 
    #                                             groupby=groupby)
    # df_res[['pupil_lick_s', 'pupil_lick_s_c']] = (df_res[['pupil_lick_s', 'pupil_lick_s_c']] +
    #                                             analysis_tools.compute_metrics(df=df.loc[df['reward_ind'] & (df['fa']==1),:], 
    #                                                                             df_cox=df_cox.loc[df_cox['reward_ind'] & (df_cox['fa']==1),:], 
    #                                                                             groupby=groupby)[['pupil_lick_s', 'pupil_lick_s_c']]) / 2
    # for y in ['pupil_lick_s', 'pupil_lick_s_c']:
    #     fig = analysis_tools.plot_scalars_across_blocks(df=df_res, x='block_id', y=y)
    #     fig.savefig(os.path.join(fig_dir, '{}_sdt_across_blocks_{}.pdf'.format(exp,y)))





    # means = df.groupby(['subject_id', 'pupil_trial_start_bin'])['pupil_lick_s'].mean().groupby(['pupil_trial_start_bin']).mean()
    # sems = df.groupby(['subject_id', 'pupil_trial_start_bin'])['pupil_lick_s'].mean().groupby(['pupil_trial_start_bin']).sem()
    # x = df.groupby(['subject_id', 'pupil_trial_start_bin'])['pupil_trial_start_c1'].mean().groupby(['pupil_trial_start_bin']).mean()

    # fig = plt.figure(figsize=(1.75, 1.75))
    # plt.errorbar(x, means, yerr=sems, fmt='-o', color='grey', markerfacecolor='lightgrey', markeredgewidth=0.75)
    # plt.xlabel('Pupil size (% max)')
    # plt.ylabel('Pupil response (% max / s)')
    # sns.despine(trim=False)
    # plt.tight_layout()
    # fig.savefig(os.path.join(fig_dir, '{}_phasic_vs_baseline_pupil.pdf'.format(exp)))

#################################
# across subject correlations:
#################################
if plot_individual_differences:

    import statsmodels.formula.api as smf

    # groupby = ['subject_id', 'session_id']
    groupby = ['subject_id']
    columns = ['hit', 'tr_d', 'tr_c', 'rt2', 'dist_from_optimal', 'pupil_trial_start_c3']

    df_res_high = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'] & (df['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),:], 
                                                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)
    # 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'])

    # formula = 'd_ps_d ~ d_ps'
    # res = smf.ols(formula=formula, data=df_res).fit()
    # df_res['d_ps_d_r'] = res.resid

    plt_nr = 1
    fig = plt.figure(figsize=(2*1.75, 1.75))
    for x, y in zip(['tr_c', 'tr_c_d'], ['tr_d', 'tr_d_d']):
        ax = fig.add_subplot(1,2,plt_nr)
        # f, p, _, _, order = analysis_tools.sequential_regression(df_res, x, y, order=1)
        r,p = sp.stats.pearsonr(df_res[x], df_res[y])
        sns.regplot(x=df_res[x], y=df_res[y], order=1, ax=ax)
        plt.title('r={}, p={}'.format(round(r,2), round(p,3)))
        plt_nr += 1
    plt.tight_layout()
    fig.savefig(os.path.join(fig_dir, '{}_across_correlations_1.pdf'.format(exp)))


    df_res['tr_c_d'] = df_res['tr_c_d'] * -1
    plt_nr = 1
    fig = plt.figure(figsize=(2*1.75, 1.75))
    for x, y in zip(['tr_d_d', 'tr_c_d'], ['hit_d', 'hit_d']):
        ax = fig.add_subplot(1,2,plt_nr)
        # f, p, _, _, order = analysis_tools.sequential_regression(df_res, x, y, order=1)
        r,p = sp.stats.pearsonr(df_res[x], df_res[y])
        sns.regplot(x=df_res[x], y=df_res[y], order=1, ax=ax)
        plt.title('r={}, p={}'.format(round(r,2), round(p,3)))
        plt_nr += 1
    plt.tight_layout()
    fig.savefig(os.path.join(fig_dir, '{}_across_correlations_2.pdf'.format(exp)))
    print(analysis_tools.permutationTest_correlation(df_res[['tr_d_d', 'hit_d']].values, 
                                                     df_res[['tr_c_d', 'hit_d']].values))

    plt_nr = 1
    fig = plt.figure(figsize=(2*1.75, 1.75))
    for x, y in zip(['tr_d', 'tr_c'], ['tr_d_d', 'tr_d_d']):
        ax = fig.add_subplot(1,2,plt_nr)
        # f, p, _, _, order = analysis_tools.sequential_regression(df_res, x, y, order=1)
        r,p = sp.stats.pearsonr(df_res[x], df_res[y])
        sns.regplot(x=df_res[x], y=df_res[y], order=1, ax=ax)
        plt.title('r={}, p={}'.format(round(r,2), round(p,3)))
        plt_nr += 1
    plt.tight_layout()
    fig.savefig(os.path.join(fig_dir, '{}_across_correlations_3.pdf'.format(exp)))


    # plt_nr = 1
    # fig = plt.figure(figsize=(2*1.75,1.75))
    # for x, y in zip(['tr_d', 'tr_c'], ['tr_d_d', 'tr_c_d']):
    #     ax = fig.add_subplot(1,2,plt_nr)
    #     f, p, _, _, order = analysis_tools.sequential_regression(df_res, x, y, order=1)
    #     sns.regplot(x=df_res[x], y=df_res[y], order=order, ax=ax)
    #     formula = '{} ~ {}'.format(y,x)
    #     if order == 2:
    #         formula = formula + ' + I({})'.format(x)
    #     model = smf.ols(formula=formula, data=df_res)
    #     res = model.fit()
    #     plt.title('r={}, p={}'.format(round(np.sqrt(res.rsquared),2), round(p[order],3)))
    #     plt_nr += 1
    # plt.tight_layout()
    # fig.savefig(os.path.join(fig_dir, '{}_across_correlations_1.pdf'.format(exp)))

    # print(permutationTest_correlation(df_res[['d_ps', 'd_ps_d']].values, df_res[['r_ps', 'd_ps_d']].values))

    # formula = 'd_ps_d ~ d_ps + r_ps'
    # model = smf.ols(formula=formula, data=df_res)
    # res = model.fit()
    # print(res.summary())

    # plt_nr = 1
    # fig = plt.figure(figsize=(2*1.75,1.75))
    # for x, y in zip(['d_ps', 'r_ps'], ['hit', 'hit',]):
    #     ax = fig.add_subplot(1,2,plt_nr)
    #     f, p, _, _, order = analysis_tools.sequential_regression(df_res, x, y, order=1)
    #     sns.regplot(df_res[x], df_res[y], order=order, ax=ax)
    #     formula = '{} ~ {}'.format(y,x)
    #     if order == 2:
    #         formula = formula + ' + I({})'.format(x)
    #     model = smf.ols(formula=formula, data=df_res)
    #     res = model.fit()
    #     plt.title('r={}, p={}'.format(round(np.sqrt(res.rsquared),2), round(p[order],3)))
    #     plt_nr += 1
    # plt.tight_layout()
    # fig.savefig(os.path.join(fig_dir, '{}_across_correlations_2.pdf'.format(exp)))

    # print(permutationTest_correlation(df_res[['d_ps', 'hit']].values, df_res[['r_ps', 'hit']].values))

    # formula = 'hit ~ d_ps + r_ps'
    # model = smf.ols(formula=formula, data=df_res)
    # res = model.fit()
    # print(res.summary())









    # plt_nr = 1
    # fig = plt.figure(figsize=(3*1.75,1.75))
    # for x, y in zip(['r_ps', 'd_ps', 'r_ps'], ['d_ps', 'hit',  'hit']):
    #     ax = fig.add_subplot(1,3,plt_nr)
    #     f, p, _, _, order = analysis_tools.sequential_regression(df_res, x, y, order=2)
    #     sns.regplot(df_res[x], df_res[y], order=order, ax=ax)
    #     formula = '{} ~ {}'.format(y,x)
    #     if order == 2:
    #         formula = formula + ' + I({})'.format(x)
    #     model = smf.ols(formula=formula, data=df_res)
    #     res = model.fit()
    #     plt.title('r={}, p={}'.format(round(np.sqrt(res.rsquared),2), round(p[order],3)))
    #     plt_nr += 1
    # plt.tight_layout()
    # fig.savefig(os.path.join(fig_dir, '{}_across_correlations_2.pdf'.format(exp)))

    # plt_nr = 1
    # fig = plt.figure(figsize=(3*1.75,1.75))
    # for x, y in zip(['r_ps_d', 'd_ps_d', 'r_ps_d'], ['d_ps_d', 'hit_d',  'hit_d']):
    #     ax = fig.add_subplot(1,3,plt_nr)
    #     f, p, _, _, order = analysis_tools.sequential_regression(df_res, x, y, order=2)
    #     sns.regplot(df_res[x], df_res[y], order=order, ax=ax)
    #     formula = '{} ~ {}'.format(y,x)
    #     if order == 2:
    #         formula = formula + ' + I({})'.format(x)
    #     model = smf.ols(formula=formula, data=df_res)
    #     res = model.fit()
    #     plt.title('r={}, p={}'.format(round(np.sqrt(res.rsquared),2), round(p[order],3)))
    #     plt_nr += 1
    # plt.tight_layout()
    # fig.savefig(os.path.join(fig_dir, '{}_across_correlations_3.pdf'.format(exp)))


    # formula = 'hit ~ d_ps + I(d_ps**2) + r_ps'
    # model = smf.ols(formula=formula, data=df_res)
    # res = model.fit()
    # print(res.summary())

    # formula = 'hit_d ~ d_ps_d + I(d_ps_d**2) + r_ps_d + I(r_ps_d**2)'
    # model = smf.ols(formula=formula, data=df_res)
    # res = model.fit()
    # print(res.summary())


    # counts = df.groupby(['subject_id']).count()
    # sns.regplot(x=counts['trial'], y=df_res['d_ps'])
    # r,p = sp.stats.pearsonr(x=counts['trial'], y=df_res['d_ps'])
    # print(r)


    # plt_nr = 1
    # fig = plt.figure(figsize=(len(columns)*1.75,len(columns)*1.75))
    # for x in columns:
    #     for y in columns:
    #         ax = fig.add_subplot(len(columns),len(columns),plt_nr)
    #         sns.regplot(df_res_mean[x], df_res_mean[y], ax=ax)
    #         r,p = sp.stats.pearsonr(df_res_mean[x], df_res_mean[y])
    #         plt.title('r={}, p={}'.format(round(r,2), round(p,3)))
    #         plt_nr += 1
    # plt.tight_layout()
    # fig.savefig(os.path.join(fig_dir, '{}_across_correlations_mean_mean.pdf'.format(exp)))


    # plt_nr = 1
    # fig = plt.figure(figsize=(len(columns)*1.75,len(columns)*1.75))
    # for x in columns:
    #     for y in columns:
    #         ax = fig.add_subplot(len(columns),len(columns),plt_nr)
    #         sns.regplot(df_res_diff[x], df_res_diff[y], ax=ax)
    #         r,p = sp.stats.pearsonr(df_res_diff[x], df_res_diff[y])
    #         plt.title('r={}, p={}'.format(round(r,2), round(p,3)))
    #         plt_nr += 1
    # plt.tight_layout()
    # fig.savefig(os.path.join(fig_dir, '{}_across_correlations_diff_diff.pdf'.format(exp)))


    # plt_nr = 1
    # fig = plt.figure(figsize=(len(columns)*1.75,len(columns)*1.75))
    # for x in columns:
    #     for y in columns:
    #         ax = fig.add_subplot(len(columns),len(columns),plt_nr)
    #         sns.regplot(df_res_mean[x], df_res_diff[y], ax=ax)
    #         r,p = sp.stats.pearsonr(df_res_mean[x], df_res_diff[y])
    #         plt.title('r={}, p={}'.format(round(r,2), round(p,3)))
    #         plt_nr += 1
    # plt.tight_layout()
    # fig.savefig(os.path.join(fig_dir, '{}_across_correlations_mean_diff.pdf'.format(exp)))

#################################
# pupil behavior correlations:
#################################
if plot_pupil_behavior_correlations:

    # add pupil bins:
    q = 8
    ind = (df['block_id']>0)&(df['reward_ind'])
    df.loc[ind&(df['walk']==0), 'pupil_trial_start_c1_bin_q'] = df.loc[ind&(df['walk']==0),:].groupby(['subject_id'], group_keys=False)['pupil_trial_start_c1'].apply(pd.qcut, q=q, labels=False).astype(int)
    df.loc[ind&(df['walk']==1), 'pupil_trial_start_c1_bin_q'] = 100
    # df.loc[ind&(df['walk']==0), 'pupil_trial_start_c1_bin_q2'] = df.loc[ind&(df['walk']==0),:].groupby(['subject_id', 'block_id'], group_keys=False)['pupil_trial_start_c1'].apply(pd.qcut, q=q, labels=False).astype(int)
    # df.loc[ind&(df['walk']==1), 'pupil_trial_start_c1_bin_q2'] = 100


    # df.loc[(df['block_id']>0)&(df['walk']==0), 'pupil_trial_start_bin_q'] = df.loc[(df['block_id']>0)&(df['walk']==0),:].groupby(['subject_id', 'reward'], group_keys=False)['pupil_trial_start'].apply(pd.qcut, q=8, labels=False).astype(int)
    # df.loc[(df['block_id']>0)&(df['walk']==1), 'pupil_trial_start_bin_q'] = 100
    # df.loc[(df['block_id']>0)&(df['walk']==0), 'pupil_trial_start_c1_bin_q'] = df.loc[(df['block_id']>0)&(df['walk']==0),:].groupby(['subject_id', 'reward'], group_keys=False)['pupil_trial_start_c1'].apply(pd.qcut, q=8, labels=False).astype(int)
    # df.loc[(df['block_id']>0)&(df['walk']==1), 'pupil_trial_start_c1_bin_q'] = 100
    # df.loc[(df['block_id']>0)&(df['walk']==0), 'pupil_trial_start_c2_bin_q'] = df.loc[(df['block_id']>0)&(df['walk']==0),:].groupby(['subject_id', 'reward'], group_keys=False)['pupil_trial_start_c2'].apply(pd.qcut, q=8, labels=False).astype(int)
    # df.loc[(df['block_id']>0)&(df['walk']==1), 'pupil_trial_start_c2_bin_q'] = 100

    for bin_by in ['pupil_trial_start_bin', 'pupil_trial_start_c3_bin']:
        # , 'pupil_trial_start_c3_bin']: # 'pupil_trial_start_bin', 'pupil_trial_start_c2_bin', 'pupil_trial_start_c3_bin']:
    # for bin_by in ['pupil_trial_start_c1_bin']: 
        
        import imp
        imp.reload(analysis_tools)

        groupby = ['subject_id', 'reward', bin_by]
        df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'],:], groupby=groupby)
        fig, axes = plt.subplots(nrows=1, ncols=5, figsize=(5*1.5,1.75))
        for r in [0,1]:
            axes = analysis_tools.baseline_pupil_behavior(df_res=df_res.loc[df_res['reward']==r,:], x_measure=bin_by[:-4], 
                                                        bin_by=bin_by, color=sns.color_palette()[r], axes=axes)
        fig.savefig(os.path.join(fig_dir, '{}_{}_{}_behavior.pdf'.format(exp, 'pupil', bin_by)))

        groupby = ['subject_id', bin_by]
        df_res = df_res.groupby(groupby).mean().reset_index()
        
        fig = analysis_tools.baseline_pupil_behavior(df_res=df_res, x_measure=bin_by[:-4], 
                                                     bin_by=bin_by, color='grey', axes=None)
        fig.savefig(os.path.join(fig_dir, '{}_{}_{}_behavior_all.pdf'.format(exp, 'pupil', bin_by)))


        groupby = ['subject_id', bin_by]
        df_res = df_res.groupby(groupby).mean().reset_index()
        
        fig = analysis_tools.baseline_pupil_behavior(df_res=df_res, x_measure=bin_by[:-4],
                                                     measures = ['tr_hr', 'tr_far', 'reward_ml', 'reward_ml', 'reward_ml'],
                                                     bin_by=bin_by, color='grey', axes=None)
        fig.savefig(os.path.join(fig_dir, '{}_{}_{}_behavior_all_2.pdf'.format(exp, 'pupil', bin_by)))



    # CONTROL FOR EARLY AND LATE WITHIN BLOCK:
    for bin_by in ['pupil_trial_start_bin', 'pupil_trial_start_c3_bin']:
        fig, axes = plt.subplots(nrows=1, ncols=5, figsize=(5*1.5,1.75))
        groupby = ['subject_id', 'reward', bin_by]
        for r in [0,1]:
            if r == 0:
                df_res = analysis_tools.compute_metrics(df=df.loc[(df['trial_since_hit']>1)&(df['trial_since_hit']<25),:], groupby=groupby)
            elif r == 1:
                df_res = analysis_tools.compute_metrics(df=df.loc[(df['trial_since_hit']>25),:], groupby=groupby)
            groupby = ['subject_id', bin_by]
            df_res = df_res.groupby(groupby).mean().reset_index()
            axes = analysis_tools.baseline_pupil_behavior(df_res=df_res, x_measure=bin_by[:-4], 
                                                        bin_by=bin_by, color=sns.color_palette()[r+2], axes=axes)
        fig.savefig(os.path.join(fig_dir, '{}_{}_{}_behavior_all_block_earlylate.pdf'.format(exp, 'pupil', bin_by)))

    # CONTROL FOR EARLY AND LATE SESSIONS:
    # ------------------------------------

    df_all = pd.read_csv(os.path.join(data_dir, 'data_learning.csv'))
    df_res = df_all.groupby(['subject_id', 'session_id'])['phase'].max().reset_index()
    dfs = []
    for phase, d_phase in df_res.groupby(['phase']):
        
        # allign:
        dfs_ = []
        for subj, d_subj in d_phase.groupby(['subject_id']):

            # print(phase)
            
            # remove weird solo sessions:
            if phase[0] == 3:
                if (d_subj['session_id'].diff().iloc[0:10]>2).sum() > 1:
                    print()
                    print('hello')
                    print(d_subj)
            if phase[0] == 2:
                print('now!')
                d_subj['session_id_phase'] = d_subj['session_id'] - d_subj['session_id'].max()
            else:
                d_subj['session_id_phase'] = d_subj['session_id'] - d_subj['session_id'].min()
            dfs_.append(d_subj)
        d_phase = pd.concat(dfs_)

        # append:
        dfs.append(d_phase)
    df_res_ = pd.concat(dfs)
    df_res_ = df_res_.loc[df_res_['phase']==3,:]
    df = df.merge(df_res_, on=['subject_id', 'session_id'])

    # groupby = ['subject_id', 'session_id_phase']
    # df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind']&(df['session_id_phase']<50),:], groupby=groupby)

    # df_res.groupby(['session_id_phase'])['tr_d'].mean()

    for bin_by in ['pupil_trial_start_bin', 'pupil_trial_start_c3_bin']:
        
        import imp
        imp.reload(analysis_tools)

        max_session = df.groupby(['subject_id'])['session_id_phase'].max().reset_index()
        max_session = max_session.loc[max_session['session_id_phase']>=50,:]
        print('nr_subjects:')
        print(max_session.shape)

        fig, axes = plt.subplots(nrows=1, ncols=5, figsize=(5*1.5,1.75))
        groupby = ['subject_id', 'reward', bin_by]
        for r in [0,1]:
            if r == 0:
                df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'] 
                                                               & (df['session_id_phase']<20)
                                                               & (df['session_id_phase']<50)
                                                               & (df['subject_id'].isin(max_session['subject_id'])),:], groupby=groupby)
            elif r == 1:
                df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'] 
                                                               & (df['session_id_phase']>30)
                                                               & (df['session_id_phase']<50)
                                                               & (df['subject_id'].isin(max_session['subject_id'])),:], groupby=groupby)
            groupby = ['subject_id', bin_by]
            df_res = df_res.groupby(groupby).mean().reset_index()
            axes = analysis_tools.baseline_pupil_behavior(df_res=df_res, x_measure=bin_by[:-4], 
                                                        bin_by=bin_by, color=sns.color_palette()[r+2], axes=axes)
        fig.savefig(os.path.join(fig_dir, '{}_{}_{}_behavior_all_session_earlylate.pdf'.format(exp, 'pupil', bin_by)))

#################################
# Cox regression:
#################################
if plot_cox_regression:

    # prepare data:
    n_p_trials = 7
    df['trial_dur'] = (df['time_trial_end']-df['time_trial_start'])
    df_cox = analysis_tools.cox_prepare_data(df)
    df_cox = df_cox.loc[df_cox['reward_ind']]
    
    # define models:
    m0_predictors = '1 + trial + reward + signal + signal:trial + signal:reward' 
    previous_predictors = ' + '.join(['reward_obtained_p{} + signal:reward_obtained_p{}'.format(p,p) for p in range(1,n_p_trials+1)])
    walk_predictors = '{} + signal:{}'.format('walk', 'walk')
    walk_p_predictors = '{} + signal:{}'.format('walk_p', 'walk_p')
    pupil_predictors = '{} + signal:{}'.format('pupil_trial_start', 'pupil_trial_start')
    pupil2_predictors = '{} + signal:{}'.format('I(pupil_trial_start**2)', 'I(pupil_trial_start**2)')
    dist_predictors = '{} + signal:{}'.format('dist_from_optimal', 'dist_from_optimal')
    # previous_predictors2 = ' + '.join(['reward_obtained2_p{} + signal:reward_obtained2_p{}'.format(p,p) for p in range(1,n_p_trials)])
    # previous_predictors3 = ' + '.join(['hit_p{} + hit_p{}:reward + signal:hit_p{} + signal:hit_p{}:reward'.format(p,p,p,p) for p in range(1,n_p_trials)])
    # reward_diff_predictors = 'reward_obtained_p1_diff + signal:reward_obtained_p1_diff'
    local_reward_rate_predictors = 'local_reward_rate_c + signal:local_reward_rate_d'

    m0 = '{}'.format(m0_predictors)
    m1 = '{} + {}'.format(m0_predictors, previous_predictors)
    m2 = '{} + {} + {} + {}'.format(m0_predictors, walk_predictors, pupil_predictors, pupil2_predictors)
    m3 = '{} + {} + {} + {} + {}'.format(m0_predictors, previous_predictors, walk_predictors, pupil_predictors, pupil2_predictors)
    m4 = '{} + {} + {} + {} + {} + {}'.format(m0_predictors, previous_predictors, walk_predictors, pupil_predictors, pupil2_predictors, walk_p_predictors)

    # m4 = '{} + {} + {}'.format(m0_predictors, previous_predictors, dist_predictors)

    # fit models:
    groupby = 'subject_id'
    df_res = []
    for m in [m0, m1, m2, m3, m4]:
        print(m)
        res = analysis_tools.cox_regression_group(df=df_cox, formula=m, groupby=[groupby], n_jobs=8, effect_coding=True)
        df_res.append(res)

    # plot:
    orders = [['trial', 'reward', 'signal:trial', 'signal:reward'],
            
            ['trial', 'reward'] + 
             ['reward_obtained_p{}'.format(p) for p in range(1,n_p_trials+1)] +
             ['signal:trial', 'signal:reward'] + 
             ['signal:reward_obtained_p{}'.format(p) for p in range(1,n_p_trials+1)],
            
            ['trial', 'reward'] + 
             ['walk', 'pupil_trial_start', 'I(pupil_trial_start**2)'] + 
             ['signal:trial', 'signal:reward'] + 
             ['signal:walk', 'signal:pupil_trial_start', 'signal:I(pupil_trial_start**2)'],

            ['trial', 'reward'] + 
             ['reward_obtained_p{}'.format(p) for p in range(1,n_p_trials+1)] +
             ['walk', 'pupil_trial_start', 'I(pupil_trial_start**2)'] + 
             ['signal:trial', 'signal:reward'] + 
             ['signal:reward_obtained_p{}'.format(p) for p in range(1,n_p_trials+1)] +
             ['signal:walk', 'signal:pupil_trial_start', 'signal:I(pupil_trial_start**2)'],

            ['trial', 'reward'] + 
             ['reward_obtained_p{}'.format(p) for p in range(1,n_p_trials+1)] +
             ['walk', 'walk_p', 'pupil_trial_start', 'I(pupil_trial_start**2)'] + 
             ['signal:trial', 'signal:reward'] + 
             ['signal:reward_obtained_p{}'.format(p) for p in range(1,n_p_trials+1)] +
             ['signal:walk', 'signal:walk_p', 'signal:pupil_trial_start', 'signal:I(pupil_trial_start**2)'],

            # ['trial', 'reward'] + 
            #  ['reward_obtained_p{}'.format(p) for p in range(1,n_p_trials+1)] +
            #  ['dist_from_optimal'] + 
            #  ['signal:trial', 'signal:reward'] + 
            #  ['signal:reward_obtained_p{}'.format(p) for p in range(1,n_p_trials+1)] +
            #  ['signal:dist_from_optimal'],
            ]


    names = [['trial', 'reward', 'trial', 'reward'],
            
            ['trial', 'reward'] + 
             ['outcome_p{}'.format(p) for p in range(1,n_p_trials+1)] +
             ['trial', 'reward'] + 
             ['outcome_p{}'.format(p) for p in range(1,n_p_trials+1)],

            ['trial', 'reward'] + 
             ['walk', 'pupil', 'pupil^2',] +
             ['trial', 'reward'] + 
             ['walk', 'pupil', 'pupil^2',],

            ['trial', 'reward'] + 
             ['outcome_p{}'.format(p) for p in range(1,n_p_trials+1)] +
             ['walk', 'pupil', 'pupil^2',] +
             ['trial', 'reward'] + 
             ['outcome_p{}'.format(p) for p in range(1,n_p_trials+1)] +
             ['walk', 'pupil', 'pupil^2',],

            ['trial', 'reward'] + 
             ['outcome_p{}'.format(p) for p in range(1,n_p_trials+1)] +
             ['walk', 'walk_p', 'pupil', 'pupil^2',] +
             ['trial', 'reward'] + 
             ['outcome_p{}'.format(p) for p in range(1,n_p_trials+1)] +
             ['walk', 'walk_p', 'pupil', 'pupil^2',],

            # ['trial', 'reward'] + 
            #  ['outcome_p{}'.format(p) for p in range(1,n_p_trials+1)] +
            #  ['dist',] +
            #  ['trial', 'reward'] + 
            #  ['outcome_p{}'.format(p) for p in range(1,n_p_trials+1)] +
            #  ['dist',],
             ]

    for m in [0,1,2,3,4]:

        print(df_res[m]['ll'].mean())
        print(df_res[m]['aic'].mean())
        # print(df_res[m].columns)

        df_plot = df_res[m].drop(['index'], axis=1).melt(id_vars=groupby)
        df_plot['value'] = np.log(df_plot['value'])

        # fig = plt.figure(figsize=(1+len(orders[m])/4,1.75))
        fig = plt.figure(figsize=(1.5+len(orders[m])/4,3))
        # gs = fig3.add_gridspec(1, 6)
        for plot_nr, title in enumerate(['bias', 'sens']):
            ax = fig.add_subplot(1,2,plot_nr+1)
            if plot_nr == 0:
                variables = orders[m][:int(len(orders[m])/2)]
                variable_names = names[m][:int(len(orders[m])/2)]
            elif plot_nr == 1:
                variables = orders[m][int(len(orders[m])/2):]
            sns.barplot(x='covariate', y='value', order=variables, 
                        linewidth=0.5, 
                        data=df_plot, ax=ax)
            plt.axhline(0, color='k', lw=1, ls='--')
            plt.title(title)
            plt.ylabel('Hazard ratio')
            import matplotlib.transforms as transforms
            trans = transforms.blended_transform_factory(ax.transData, ax.transAxes)
            for i, v in enumerate(variables):
                # t,p = sp.stats.wilcoxon(df_lr.loc[(df_lr['variable']==v)&(df_lr['reward']==1), 'value'], df_lr.loc[(df_lr['variable']==v)&(df_lr['reward']==0), 'value'])
                t,p = sp.stats.ttest_1samp(df_plot.loc[(df_plot['covariate']==v), 'value'], 0)
                plt.text(x=i, y=0.9, s='p={}'.format(round(p,3)), size=6, rotation=45, transform=trans)
            plt.xticks(ticks=np.arange(len(variables)), labels=variable_names, rotation=45)
            plt.xlabel('')
        sns.despine(trim=False)
        plt.tight_layout()
        fig.savefig(os.path.join(fig_dir, '{}_cox_regression_{}.pdf'.format(exp, m)))

    
    # add local utility:
    # df_res = analysis_tools.cox_regression_group(df=df_cox, formula=m1, groupby=['subject_id'], n_jobs=8, effect_coding=True)
    # df = df.groupby(['subject_id']).apply(compute_local_utility, 
    #                                  df_res=df_res, 
    #                                  name='coefs_trial',
    #                                  columns=['trial']).reset_index(drop=True)

    # df = df.groupby(['subject_id']).apply(compute_local_utility, 
    #                                  df_res=df_res, 
    #                                  name='coefs_reward',
    #                                  columns=['reward']).reset_index(drop=True)

    df = df.groupby(['subject_id']).apply(compute_local_utility, 
                                     df_res=df_res[1], 
                                     name='coefs_history',
                                     columns=['reward_obtained_p{}'.format(p) 
                                                    for p in range(1,n_p_trials)]).reset_index(drop=True)

    df = df.groupby(['subject_id']).apply(compute_local_utility, 
                                     df_res=df_res[1], 
                                     name='coefs_all',
                                     columns=['trial', 'reward']+['reward_obtained_p{}'.format(p) 
                                                                  for p in range(1,n_p_trials)]).reset_index(drop=True)






    # plt.plot(df.groupby(['trial'])['local_utility_c'].mean(), label='bias')
    # plt.plot(df.groupby(['trial'])['local_utility_d'].mean(), label='sens')
    # plt.plot(df.groupby(['trial'])['local_utility'].mean(), label='all')
    # plt.legend()

    # ind = ~df['dist_from_optimal'].isna()
    # sp.stats.pearsonr(df.loc[ind,'dist_from_optimal'], df.loc[ind,'local_utility_c'])

    # dfs = []
    # for i in range(len(df_res)):
    #     d = df_res[i][['subject_id', 'signal:reward']]
    #     d['model'] = i
    #     dfs.append(d)
    # df_reward = pd.concat(dfs)

#################################
# mediation analysis:
#################################
if plot_mediation:

    # # single trial lavaan:

    # columns = ['subject_id', 'session_id', 'block_id', 'reward', 'coefs_trial', 'coefs_reward', 'coefs_history', 'coefs_all',
    #                         'dist_from_optimal', 'pupil_trial_start_c3', 'pupil_trial_start', 'walk',
    #                         'hit']
    # df_r = []
    # for (subj,), d in df.loc[df['reward_ind'], columns].groupby(['subject_id']):
    #     print(subj)
    #     df_r.append(d)
    # df_r = pd.concat(df_r)
    # # df_r = df_r.groupby(['subject_id', 'session_id', 'block_id']).mean().reset_index()
    # df_r.to_csv(os.path.join(fig_dir, 'data_r.csv'))

    # for m in ['dist_from_optimal', 'pupil_trial_start_c3', 'walk']:
    #     # m = 'dist_from_optimal'
    #     dfs = []
    #     for x in ['reward', 'coefs_history', 'coefs_all']:
    #         # x = 'reward'
    #         data = pd.read_csv(os.path.join(fig_dir, 'lavaan_results_{}_{}.csv'.format(x,m)))
    #         data.loc[data['label']=='indirect', 'est'] = data.loc[data['label']=='indirect', 'est'] / data.loc[data['label']=='total', 'est'].values.mean() * 100        
    #         data = data.loc[data['label']=='indirect', :]
    #         data['x'] = x
    #         print(sp.stats.ttest_1samp(data['est'],0))
    #         dfs.append(data)
    #     data_med = pd.concat(dfs)
        
    #     fig = plt.figure(figsize=(2,2))
    #     ax = fig.add_subplot(111)
    #     sns.barplot(x='x', y='est', errorbar='se', data=data_med, ax=ax)
    #     import matplotlib.transforms as transforms
    #     trans = transforms.blended_transform_factory(ax.transData, ax.transAxes)
    #     p1 = sp.stats.ttest_1samp(data_med.loc[data_med['x']=='reward', 'est'], 0)[1]
    #     p2 = sp.stats.ttest_1samp(data_med.loc[data_med['x']=='coefs_history', 'est'], 0)[1]
    #     p3 = sp.stats.ttest_1samp(data_med.loc[data_med['x']=='coefs_all', 'est'], 0)[1]
    #     plt.text(x=0, y=0.75, s='p={}'.format(round(p1,3)), size=6, rotation=45, transform=trans)
    #     plt.text(x=1, y=0.75, s='p={}'.format(round(p2,3)), size=6, rotation=45, transform=trans)
    #     plt.text(x=2, y=0.75, s='p={}'.format(round(p3,3)), size=6, rotation=45, transform=trans)
    #     plt.ylabel('Coefficient (a.u.)')
    #     ax.legend([],[], frameon=False)
    #     sns.despine(trim=False)
    #     plt.tight_layout()
    #     fig.savefig(os.path.join(fig_dir, '{}_mediation_{}_{}_{}_{}.pdf'.format(exp, x, m, path, walk)))

    for walk in ['all', 'still']:

        # logistic regression:
        groupby = ['subject_id', 'session_id', 'block_id']
        if walk == 'all':
            df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'],:], 
                                                        groupby=groupby)
            counts = df.loc[df['reward_ind'],:].groupby(groupby)['trial'].count().reset_index()
        elif walk == 'still':
            df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'] & (df['walk']==0) ,:], 
                                                        groupby=groupby)
            counts = df.loc[df['reward_ind'] & (df['walk']==0),:].groupby(groupby)['trial'].count().reset_index()
        df_res['reward'] = ((df_res['block_id']==1)|(df_res['block_id']==3)|(df_res['block_id']==5)).astype(int)
        
        df_res = df_res.loc[np.array(counts['trial']>=20),:]

        # mediation:
        from pingouin import mediation_analysis
        for x in ['reward', 'coefs_history', 'coefs_all']:
        # for x in ['reward']:
            for m in ['dist_from_optimal', 'pupil_trial_start', 'pupil_trial_start_c3']: # 'pupil_trial_start_c1_std',
            # for m in ['dist_from_optimal']:

                dfs = []
                for i, y in enumerate(['tr_c', 'tr_d', 'hit']):
                    for (subj,), data in df_res.groupby(['subject_id']):
                        data[x] = (data[x]-data[x].mean()) / data[x].std()
                        data[m] = (data[m]-data[m].mean()) / data[m].std()
                        data[y] = (data[y]-data[y].mean()) / data[y].std()
                        res = mediation_analysis(data=data, x=x, m=m, y=y, alpha=0.05, n_boot=1)
                        res['subject_id'] = subj
                        res['y'] = y
                        dfs.append(res)
                res = pd.concat(dfs)

                # res.loc[res['path']=='Indirect','coef'] = res.loc[res['path']=='Indirect','coef'] / res.loc[res['path']=='Total','coef'].mean()
                # res = res.loc[res['path']=='Indirect',:]
                # print(res['coef'].mean()*100)
                

                # for path in ['Total', 'Direct', 'Indirect']:
                for path in ['Indirect']:
                    
                    coefs = res.loc[res['path']==path, ['y', 'coef', 'subject_id']]
                    coefs.loc[coefs['y']=='tr_c', 'coef'] = coefs.loc[coefs['y']=='tr_c', 'coef']*-1


                    p1 = sp.stats.wilcoxon(coefs.loc[coefs['y']=='tr_c', 'coef'], np.zeros(coefs.loc[coefs['y']=='tr_d', 'coef'].shape[0]))[1]
                    p2 = sp.stats.wilcoxon(coefs.loc[coefs['y']=='tr_d', 'coef'], np.zeros(coefs.loc[coefs['y']=='tr_c', 'coef'].shape[0]))[1]
                    p3 = sp.stats.wilcoxon(coefs.loc[coefs['y']=='hit', 'coef'], np.zeros(coefs.loc[coefs['y']=='hit', 'coef'].shape[0]))[1]
                    p4 = sp.stats.wilcoxon(coefs.loc[coefs['y']=='tr_d', 'coef'], coefs.loc[coefs['y']=='tr_c', 'coef'])[1]

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

                    # plt.errorbar(x=[0,1,2], y=coefs.groupby(['y'])['coef'].mean().iloc[[1,2,0]], 
                    #             yerr=coefs.groupby(['y'])['coef'].sem().iloc[[1,2,0]],
                    #             fmt='o')
                    plt.xticks([0,1,2],['c', 'd', 'h'])
                    sns.barplot(x='y', y='coef', errorbar='se', data=coefs, ax=ax)

                    import matplotlib.transforms as transforms
                    trans = transforms.blended_transform_factory(ax.transData, ax.transAxes)
                    plt.text(x=0, y=0.75, s='p={}'.format(round(p1,3)), size=6, rotation=0, transform=trans)
                    plt.text(x=1, y=0.75, s='p={}'.format(round(p2,3)), size=6, rotation=0, transform=trans)
                    plt.text(x=2, y=0.75, s='p={}'.format(round(p3,3)), size=6, rotation=0, transform=trans)
                    plt.text(x=0.5, y=0.85, s='p={}'.format(round(p4,3)), size=6, rotation=0, transform=trans)
                    plt.ylabel('Coefficient (a.u.)')
                    ax.legend([],[], frameon=False)
                    sns.despine(trim=False)
                    plt.tight_layout()
                    fig.savefig(os.path.join(fig_dir, '{}_mediation_{}_{}_{}_{}.pdf'.format(exp, x, m, path, walk)))


#################################
# state occupancy analysis:
#################################

def compute_fraction_of_effect(df):

    columns = ['tr_c', 'tr_d', 'rt2', 'hit',
                'tr_c_predicted', 'tr_d_predicted', 'rt2_predicted', 'hit_predicted']

    groupby = ['subject_id', 'block_id']
    df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'],:], groupby=groupby)
    df_res['reward'] = df_res['block_id']%2
    df_res = df_res.groupby(['subject_id', 'reward']).mean().reset_index()

    df_res = (df_res.loc[df_res['reward']==1,columns].median()-
                df_res.loc[df_res['reward']==0,columns].median())
                
    # df_res = (df.loc[df['reward']==1,columns].mean()-
    #             df.loc[df['reward']==0,columns].mean())

    df_res['tr_c_predicted_f'] = df_res['tr_c_predicted']/df_res['tr_c']*100
    df_res['tr_d_predicted_f'] = df_res['tr_d_predicted']/df_res['tr_d']*100
    df_res['hit_predicted_f'] = df_res['hit_predicted']/df_res['hit']*100
    df_res['rt2_predicted_f'] = df_res['rt2_predicted']/df_res['rt2']*100

    return df_res

if plot_state_occupancy_analysis:

    # for pupil_measure in ['pupil_trial_start']:
    for pupil_measure in ['pupil_trial_start', 'pupil_trial_start_c3', 'pupil_trial_start_c3nowalk',]:
        
        bin_by = '{}_bin'.format(pupil_measure)
        
        fig = plot_delta_trial_density(df=df.loc[df['reward_ind'],:], 
                                       pupil_measure=pupil_measure)
        fig.savefig(os.path.join(fig_dir, '{}_{}_density.pdf'.format(exp, pupil_measure)))

        # predicted behavioral effect:
        df_res = analysis_tools.compute_metrics(df=df.loc[df['reward_ind'],:], groupby=['subject_id', 'reward', bin_by])
        
        # per subject:
        df_to_merge = df_res.groupby(['subject_id', bin_by])[['tr_d', 'tr_c', 'rt2', 'hit']].mean().reset_index()
        df_to_merge = df_to_merge.rename({'tr_d': 'tr_d_predicted',
                                        'tr_c': 'tr_c_predicted',
                                        'rt2': 'rt2_predicted',
                                        'hit': 'hit_predicted',}, axis=1)
        df = df.drop([c for c in df.columns if '_predicted' in c], axis=1)
        df = df.merge(df_to_merge, on=['subject_id', bin_by], how='left')

        # # across subjects:
        # df_to_merge = df_res.groupby([bin_by])[['tr_d', 'tr_c', 'rt2', 'hit']].mean().reset_index()
        # df_to_merge = df_to_merge.rename({'tr_d': 'tr_d_predicted',
        #                                   'tr_c': 'tr_c_predicted',
        #                                   'rt2': 'rt2_predicted',
        #                                   'hit': 'hit_predicted',}, axis=1)
        # df = df.drop([c for c in df.columns if '_predicted' in c], axis=1)
        # df = df.merge(df_to_merge, on=[bin_by])

        df_res2 = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(compute_fraction_of_effect)
                                                                        (df=df.groupby(['subject_id', 'block_id']).sample(frac=1, replace=True), 
                                                                        ) for i in tqdm(range(n_boot)))
        df_res2 = pd.concat(df_res2, axis=1).T
        (df_res2['tr_d_predicted_f']<0).mean()


        bars = ['tr_c_predicted_f', 'tr_d_predicted_f', 'rt2_predicted_f', 'hit_predicted_f']
        means = [df_res2[c].quantile(0.5) for c in bars]
        sems = [(df_res2[c].quantile(0.84)-df_res2[c].quantile(0.16))/2 for c in bars]
        p_values = [(df_res2[c]<0).mean() for c in bars]
        p_values.append((df_res2['tr_c_predicted_f']<df_res2['tr_d_predicted_f']).mean())
        
        fig = plt.figure(figsize=(2.5,1.75))
        ax = fig.add_subplot(111)
        ax.errorbar(x=np.arange(len(means)), y=means, yerr=sems, fmt='o')
        plt.axhline(0, color='black', lw=0.5)
        trans = transforms.blended_transform_factory(ax.transData, ax.transAxes)

        for i, p in enumerate(p_values):
            plt.text(x=i, y=0.9, s='p={}'.format(round(p,3)), size=6, rotation=45, transform=trans)
        plt.xticks(np.arange(4), ['c', 'd', 'RT (s)', 'RR'])
        plt.xlabel('')
        plt.ylabel('Effect size (%)')
        plt.tight_layout()
        sns.despine()
        fig.savefig(os.path.join(fig_dir, '{}_total_vs_predicted_{}.pdf'.format(exp,pupil_measure)))