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

import pyddm_components

# 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 = 3
n_cpus = 12
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.01
dt = 0.01
T_dur = 14.5

model_0 = {'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}

model_settings = [model_0]

# groupby = ['block_id', 'hit_p']
# version = 0
# for i in range(100)[::-1]:
#     print(i)
#     if len(groupby) == 2:
#         output_filename = os.path.join(project_dir, 'ddm_params_{}_{}_{}_sample_{}.csv'.format(version, groupby[0], groupby[1], i))
#     elif len(groupby) == 1:
#         output_filename = os.path.join(project_dir, 'ddm_params_{}_{}_sample_{}.csv'.format(version, groupby[0], i))
    
#     if not os.path.exists(output_filename):
    
#         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_resample.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)


# groupby = ['reward', 'pupil_trial_start_c3_bin']
groupby = ['reward', 'pupil_trial_start_bin']
version = 0
for i in range(100)[::-1]:
    print(i)
    if len(groupby) == 2:
        output_filename = os.path.join(project_dir, 'ddm_params_{}_{}_{}_sample_{}.csv'.format(version, groupby[0], groupby[1], i))
    elif len(groupby) == 1:
        output_filename = os.path.join(project_dir, 'ddm_params_{}_{}_sample_{}.csv'.format(version, groupby[0], i))
    
    if not os.path.exists(output_filename):
    
        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_resample.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)

    # groupby = ['reward', 'state']
    # n_cpus = 12
    # version = 5
    # for i in range(100)[::-1]:
    #     print(i)
    #     output_filename = os.path.join(project_dir, 'ddm_params_{}_{}_{}_sample_{}.csv'.format(version, groupby[0], groupby[1], i))
    #     if not os.path.exists(output_filename):
    #         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='multiprocessing')(delayed(fit_model_one_accumulator)
    #                         (data.loc[:,['start', 'duration', 'response', 'rt']], settings, n_cpus) for ids, data in df_resample.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)

    # groupby = ['reward', 'd_bin']
    # df = add_individual_difference_bins(df)
    # n_cpus = 12
    # version = 4
    # for i in range(25):
    #     print(i)
    #     output_filename = os.path.join(project_dir, 'ddm_params_{}_{}_{}_sample_{}.csv'.format(version, groupby[0], groupby[1], i))
    #     if not os.path.exists(output_filename):
    #         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='multiprocessing')(delayed(fit_model_one_accumulator)
    #                         (data.loc[:,['start', 'duration', 'response', 'rt']], settings, n_cpus) for ids, data in df_resample.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)