# Environment setup
%matplotlib inline
%cd /lang_dec
# Imports
import warnings; warnings.filterwarnings('ignore')
import hddm
import math
import scipy
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import bayesian_bootstrap.bootstrap as bootstrap
from utils import model_tools, signal_detection

# Import pilot models
pilot_data = hddm.load_csv('/lang_dec/data/pilot_clean.csv')
pilot_model = hddm.HDDM(pilot_data, depends_on={'v': 'stim'}, bias=True)
pilot_model.load_db(dbname='language_decision/models/pilot', db='txt')

#pilot_model_threshold = hddm.HDDM(pilot_data, depends_on={'v': 'stim', 'a': 'stim'})
#pilot_model_threshold.load_db(dbname='language_decision/models/pilot_threshold', db='txt')

# Import control models
controls_data = hddm.load_csv('/lang_dec/data/controls_clean.csv')
controls_model = hddm.HDDM(controls_data, depends_on={'v': 'stim'}, bias=True)
controls_model.load_db(dbname='language_decision/models/controls', db='txt')

#controls_model_threshold = hddm.HDDM(controls_data, depends_on={'v': 'stim', 'a': 'stim'}, bias=True)
#controls_model_threshold.load_db(dbname='language_decision/models/controls_threshold', db='txt')

# Import patient models
patients_data = hddm.load_csv('/lang_dec/data/patients_clean.csv')
patients_model = hddm.HDDM(patients_data, depends_on={'v': 'stim'}, bias=True)
patients_model.load_db(dbname='language_decision/models/patients', db='txt')

#patients_model_threshold = hddm.HDDM(patients_data, depends_on={'v': 'stim', 'a': 'stim'})
#patients_model_threshold.load_db(dbname='language_decision/models/patients_threshold', db='txt')

"""
Explanation: Model Comparisons
Comparing the DDM to other congitive models
End of explanation
"""


def get_d_primes(dataset, stim1, stim2, include_id=False):
    d_primes = dict()
    subject_ids = set(dataset.subj_idx)
    for subject_id in subject_ids:
        stim1_data = dataset.loc[
            dataset['subj_idx'] == subject_id].loc[
            dataset['stim'] == str(stim1)]
        stim1_trials = len(stim1_data)
        hits = len(stim1_data.loc[
            dataset['response'] == 1.0])
        
        stim2_data = dataset.loc[
            dataset['subj_idx'] == subject_id].loc[
            dataset['stim'] == str(stim2)]
        stim2_trials = len(stim2_data)
        
        fas = len(stim2_data.loc[
            dataset['response'] == 0.0])
        
        if not stim1_trials or not stim2_trials:
            d_primes[subject_id] = None # N/A placeholder value
            continue
        
        d_prime = signal_detection.signal_detection(
            n_stim1=stim1_trials, 
            n_stim2=stim2_trials, 
            hits=hits, 
            false_alarms=fas)['d_prime']
        d_primes[subject_id] = d_prime
       
    if not include_id:
        return list(d_primes.values())
    return d_primes

"""
Explanation: DDM vs Signal Detection Theory
Comparing DDM to Signal Detection - does d' correlate with DDM parameters?
End of explanation
"""


plt.hist(get_d_primes(pilot_data, 'SS', 'US'))

plt.hist(get_d_primes(pilot_data, 'SS', 'CS'))

plt.hist(get_d_primes(pilot_data, 'SS', 'CP'))

"""
Explanation: d' distributions
Pilot
End of explanation
"""


plt.hist(get_d_primes(controls_data, 'SS', 'US'))

plt.hist(get_d_primes(controls_data, 'SS', 'CS'))

plt.hist(get_d_primes(controls_data, 'SS', 'CP'))

"""
Explanation: Controls
End of explanation
"""


plt.hist(get_d_primes(patients_data, 'SS', 'US'))

plt.hist(get_d_primes(patients_data, 'SS', 'CS'))

plt.hist(list(filter(None, get_d_primes(patients_data, 'SS', 'CP'))))

"""
Explanation: Patients
End of explanation
"""


def match_dprime_to_driftrate(dataset, model, stim1, stim2):
    subject_ids = set(dataset.subj_idx)
    d_primes = get_d_primes(dataset, stim1, stim2, include_id=True)
    for subject_id in subject_ids:
        try:
            d_prime = d_primes[subject_id]
            v_stim1 = model.values['v_subj(' + stim1 + ').' + str(subject_id)]
            v_stim2 = model.values['v_subj(' + stim2 + ').' + str(subject_id)]
            v_diff = abs(v_stim2 - v_stim1)
            yield (d_prime, v_diff)
        except: 
            continue

"""
Explanation: Drift rate / d'
End of explanation
"""


dprime_driftrate = np.array([*match_dprime_to_driftrate(pilot_data, pilot_model, 'SS', 'US')])
x = dprime_driftrate[:,0]
y = dprime_driftrate[:,1]

plt.scatter(x, y)
scipy.stats.spearmanr(x, y)

dprime_driftrate = np.array([*match_dprime_to_driftrate(controls_data, controls_model, 'SS', 'US')])
x = dprime_driftrate[:,0]
y = dprime_driftrate[:,1]

plt.scatter(x, y)
scipy.stats.spearmanr(x, y)

dprime_driftrate = np.array([*match_dprime_to_driftrate(patients_data, patients_model, 'SS', 'US')])
x = dprime_driftrate[:,0]
y = dprime_driftrate[:,1]

plt.scatter(x, y)
scipy.stats.spearmanr(x, y)

"""
Explanation: SS vs US
End of explanation
"""


dprime_driftrate = np.array([*match_dprime_to_driftrate(pilot_data, pilot_model, 'SS', 'CS')])
x = dprime_driftrate[:,0]
y = dprime_driftrate[:,1]

plt.scatter(x, y)
scipy.stats.spearmanr(x, y)

dprime_driftrate = np.array([*match_dprime_to_driftrate(controls_data, controls_model, 'SS', 'CS')])
x = dprime_driftrate[:,0]
y = dprime_driftrate[:,1]

plt.scatter(x, y)
scipy.stats.spearmanr(x, y)

dprime_driftrate = np.array([*match_dprime_to_driftrate(patients_data, patients_model, 'SS', 'CS')])
x = dprime_driftrate[:,0]
y = dprime_driftrate[:,1]

plt.scatter(x, y)
scipy.stats.spearmanr(x, y)

"""
Explanation: SS vs CS
End of explanation
"""


dprime_driftrate = np.array([*match_dprime_to_driftrate(pilot_data, pilot_model, 'SS', 'CP')])
x = dprime_driftrate[:,0]
y = dprime_driftrate[:,1]

plt.scatter(x, y)
scipy.stats.spearmanr(x, y)

dprime_driftrate = np.array([*match_dprime_to_driftrate(controls_data, controls_model, 'SS', 'CP')])
x = dprime_driftrate[:,0]
y = dprime_driftrate[:,1]

plt.scatter(x, y)
scipy.stats.spearmanr(x, y)

dprime_driftrate = np.array([*match_dprime_to_driftrate(patients_data, patients_model, 'SS', 'CP')])
x = dprime_driftrate[:,0]
y = dprime_driftrate[:,1]

plt.scatter(x, y)
scipy.stats.spearmanr(x, y)

"""
Explanation: SS vs CP
End of explanation
"""


dprime_driftrate = np.array([*match_dprime_to_driftrate(patients_data, 
                                                 patients_model, 'SS', 'CS')])
x = dprime_driftrate[:,0]
y = dprime_driftrate[:,1]

plt.scatter(x, y)
scipy.stats.spearmanr(x, y)

dprime_driftrate = np.array([*match_dprime_to_driftrate(patients_data, 
                                                 patients_model, 'SS', 'CP')])
x = dprime_driftrate[:,0]
y = dprime_driftrate[:,1]

plt.scatter(x, y)
scipy.stats.spearmanr(x, y)

"""
Explanation: Low d' comparisons
Compare ddm drift rate only with low d'
Ratcliff, R. (2014). Measuring psychometric functions with the diffusion model. Journal of Experimental Psychology: Human Perception and Performance, 40(2), 870-888.
http://dx.doi.org/10.1037/a0034954
Patients are the best candidates for this (SSvsCS, SSvsCP)
End of explanation
"""


