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

import os, itertools
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")

# 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)))
# 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')
# df = pd.read_csv(os.path.join(data_dir, 'data_ddm_{}.csv'.format(exp)))
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]

    
import pyddm_components

# for version in [0,1,2,3,4,5,6]:
# for version in [5,4,3]:
for version in [0]:
    print('here we go!')
    settings = model_settings[version]
    print(settings)
    for run in range(5):
        print(version, run)
        output_filename = os.path.join(project_dir, 'ddm_params_{}_{}.csv'.format(version, run))
        # if not os.path.exists(output_filename):
        params = pyddm_components.fit_model_one_accumulator(df.loc[:,['start', 'duration', 'response', 'rt', 'reward', 'block_id', 'hit_p', 'state']], 
                                                        model_settings[version], n_cpus=n_cpus)
        params.to_csv(output_filename)

# # groupby = ['block_id']
# groupby = ['block_id', 'hit_p']
# # for version in [0,1,3,4]:
# for version in [5]:
#     if len(groupby) == 2:
#         output_filename = os.path.join(project_dir, 'ddm', 'ddm_params_{}_{}_{}.csv'.format(version, groupby[0], groupby[1]))
#     elif len(groupby) == 1:
#         output_filename = os.path.join(project_dir, 'ddm', 'ddm_params_{}_{}.csv'.format(version, groupby[0]))
#     settings = model_settings[version]
#     # df_resample = df.groupby(['subject_id'] + groupby).sample(frac=1, replace=True)
#     res = Parallel(n_jobs=n_jobs, verbose=1, backend='loky')(delayed(pyddm_components.fit_model_one_accumulator)
#                     (data.loc[:,['start', 'duration', 'response', 'rt']], settings, n_cpus) for ids, data in df.groupby(groupby))
    
#     params = pd.concat(res).reset_index(drop=True)
#     params[groupby] = pd.DataFrame([name for name, unused_df in df.groupby(groupby)], columns=groupby)
#     params.to_csv(output_filename)