'''
Purpose
How to Run
Marisa Lytle
Feb 2020
'''

import pandas as pd
import numpy as np

def deidentify(data,info):

    #Read in birthdate data
    bday_file = "./Birthdates.csv"
    birthdates = pd.read_csv(bday_file)
    birthdates['birthdate'] = pd.to_datetime(birthdates['birthdate'], format="%m/%d/%Y")

    #Convert session date to datetime object
    data['SessionDate2'] = pd.to_datetime(data['SessionDate'], format="%m-%d-%Y")

    #Merge in birthdates
    data = pd.merge(data, birthdates, on='subject', how='left')

    #Calculate age in days
    data['age'] = data['SessionDate2'] - data['birthdate']
    data['age'] = data['age'].dt.days
    
    data['ShiftedSessionDate'] = data['SessionDate2'] + pd.DateOffset(days=info['shift']) - pd.DateOffset(years=200)
    
    data = data.drop(['SessionDate','SessionDate2', 'birthdate', 'Clock.Information', 'SessionStartDateTimeUtc'], axis=1)
      
    return data

def onset_calc(run_data):
   
    #Calculate average onsets to use for subjects that cannot be calculated due to missing task or trial onsets
    if run_data['task_onset'].isnull().values.any() or run_data['trial_onset'].isnull().values.any():
        run_data['onset'] = 0
    elif run_data['task_onset'].notnull().values.any() and run_data['trial_onset'].notnull().values.any():
        run_data['onset'] = (run_data['trial_onset'] - run_data['task_onset']).round(4)

    return run_data

def clean_col(data):
    if 'Present.OnsetTime' not in data.columns:
        data = data.rename(columns={"Present.RESP": "window1_response", "Present.RT": "window1_RT",
                                  "Response.RESP": "window2_response", "Response.RT": "window2_RT", "SoundFile": "stim_file",
                                  "Condition": "trial_type", "Subject": "subject", "Response.Duration": "window2_dur", "Correct": "correct"})
        data['task_onset1'] = 0
        data['task_onset2'] = 0
        data['trial_onset'] = 0
    elif 'Blank.OnsetTime' not in data.columns:
        data = data.rename(columns={"Blank1.OnsetTime": "task_onset2",
                                  "Present.OnsetTime": "trial_onset", "Present.RESP": "window1_response", "Present.RT": "window1_RT",
                                  "Response.RESP": "window2_response", "Response.RT": "window2_RT", "SoundFile": "stim_file",
                                  "Condition": "trial_type", "Subject": "subject", "Response.Duration": "window2_dur", "Correct": "correct"})
        data['task_onset1'] = 0
    elif 'Blank1.OnsetTime' not in data.columns:
        data = data.rename(columns={"Blank.OnsetTime": "task_onset1",
                                  "Present.OnsetTime": "trial_onset", "Present.RESP": "window1_response", "Present.RT": "window1_RT",
                                  "Response.RESP": "window2_response", "Response.RT": "window2_RT", "SoundFile": "stim_file",
                                  "Condition": "trial_type", "Subject": "subject", "Response.Duration": "window2_dur", "Correct": "correct"})
        data['task_onset2'] = 0
    
    else:
        data = data.rename(columns={"Blank.OnsetTime": "task_onset1","Blank1.OnsetTime": "task_onset2",
                                  "Present.OnsetTime": "trial_onset", "Present.RESP": "window1_response", "Present.RT": "window1_RT",
                                  "Response.RESP": "window2_response", "Response.RT": "window2_RT", "SoundFile": "stim_file",
                                  "Condition": "trial_type", "Subject": "subject", "Response.Duration": "window2_dur", "Correct": "correct"})

    #initialize task onset column
    data['task_onset'] = 0

    #merge the two different task onset columns together as Gram1 uses task_onset1 and others use task_onset2
    for x in range(len(data['subject'].values)):
        if data['task_onset1'].iloc[x] > 0:
            data.loc[x,'task_onset'] = data['task_onset1'].iloc[x]
        elif data['task_onset2'].iloc[x] > 0: 
            data.loc[x,'task_onset'] = data['task_onset2'].iloc[x]
        else:
            data.loc[x,'task_onset'] = np.nan

    #Convert ms to sec
    data[['task_onset', 'trial_onset', 'window1_RT', 'window2_RT', 'window2_dur']] \
                         = data[['task_onset', 'trial_onset', 'window1_RT', 'window2_RT', 'window2_dur']].applymap(lambda x: x / 1000)

    #Take off extra characters from Experiment name and stimulus names
    data['ExperimentName'] = data['ExperimentName'].map(lambda x: x.lstrip('9_'))
    data['ExperimentName'] = data['ExperimentName'].map(lambda x: x.rstrip('_font'))
    data['ExperimentName'] = data['ExperimentName'].map(lambda x: x.rstrip('_onsets'))
    data['ExperimentName'] = data['ExperimentName'].map(lambda x: x.rstrip('_test'))
    data['ExperimentName'] = data['ExperimentName'].map(lambda x: x.rstrip(' - Final'))
    data['ExperimentName'] = data['ExperimentName'].map(lambda x: x.rstrip('_gray'))
    data['stim_file'] = data['stim_file'].map(lambda x: x.lstrip('Stimuli/'))
    
    data['window1_response_raw'] = data['window1_response']
    
    #correct reaction times listed as 0 to nan for non responses
    for x in range(len(data['window1_RT'].values)):
        if np.isnan(data['window1_response'].iloc[x]):
            data.loc[x,'window1_RT'] = np.nan
        if np.isnan(data['window2_response'].iloc[x]):
            data.loc[x,'window2_RT'] = np.nan
        if np.isnan(data['window1_RT'].iloc[x]):
            data.loc[x,'window1_response'] = np.nan 
    
    return data


def calc_rt_acc(run_data):
    control_trials = ['G_GC', 'P_SC', 'SP_Co', 'S_SC']
    #Initialize calculated columns
    run_data['calculated_response'] = 0
    run_data['calculated_RT'] = 0
    run_data['calculated_accuracy'] = 0

    run_data['duration'] = run_data["window2_dur"] + 1.75
            
    #Do the RT and ACC calculations following rules listed in ELP_TaskAccRT.docx
    for trial in range(len(run_data['subject'].values)):

        if trial == (len(run_data['subject'].values)-1):
            if pd.notnull(run_data['window2_response'].iloc[trial]) and run_data['window2_RT'].iloc[trial] > 0.25:
                run_data.loc[trial,'calculated_response'] = run_data['window2_response'].iloc[trial]
                run_data.loc[trial,'calculated_RT'] = run_data['window2_RT'].iloc[trial] + 0.75
            elif pd.isnull(run_data['window2_response'].iloc[trial]) or run_data['window2_RT'].iloc[trial] < 0.25:
                run_data.loc[trial,'calculated_response'] = np.nan
                run_data.loc[trial,'calculated_RT'] = np.nan

        elif run_data['trial_type'].iloc[trial] not in control_trials:  
            if pd.notnull(run_data['window2_response'].iloc[trial]) and run_data['window2_RT'].iloc[trial] > 0.25 :
                run_data.loc[trial,'calculated_response'] = run_data['window2_response'].iloc[trial]
                run_data.loc[trial,'calculated_RT'] = run_data['window2_RT'].iloc[trial] + 0.75
            elif (pd.isnull(run_data['window2_response'].iloc[trial]) or run_data['window2_RT'].iloc[trial] < 0.25) \
                    and (run_data['trial_type'].iloc[trial+1] in control_trials):
                run_data.loc[trial,'calculated_response'] = np.nan
                run_data.loc[trial,'calculated_RT'] = np.nan
            elif (pd.isnull(run_data['window2_response'].iloc[trial]) or run_data['window2_RT'].iloc[trial] < 0.25)\
                    and (run_data['trial_type'].iloc[trial+1] not in control_trials) and pd.notnull(run_data['window1_response'].iloc[trial+1]):
                run_data.loc[trial,'calculated_response'] = run_data['window1_response'].iloc[trial+1]
                run_data.loc[trial,'calculated_RT'] = run_data['window1_RT'].iloc[trial+1] + run_data['duration'].iloc[trial]
            elif (pd.isnull(run_data['window2_response'].iloc[trial]) or run_data['window2_RT'].iloc[trial] < 0.25)\
                 and (run_data['trial_type'].iloc[trial+1] not in control_trials) and pd.isnull(run_data['window1_response'].iloc[trial+1])\
                 and run_data['window2_RT'].iloc[trial+1] < 0.25 and pd.notnull(run_data['window2_response'].iloc[trial+1]):
                run_data.loc[trial,'calculated_response'] = run_data['window2_response'].iloc[trial+1]
                run_data.loc[trial,'calculated_RT'] = run_data['duration'].iloc[trial] + 0.75 + run_data['window2_RT'].iloc[trial+1] 
            else:
                run_data.loc[trial,'calculated_response'] = np.nan
                run_data.loc[trial,'calculated_RT'] = np.nan

        elif run_data['trial_type'].iloc[trial] in control_trials:
            if pd.notnull(run_data['window2_response'].iloc[trial]):
                run_data.loc[trial,'calculated_response'] = run_data['window2_response'].iloc[trial]
                run_data.loc[trial,'calculated_RT'] = run_data['window2_RT'].iloc[trial] + 0.75
            elif pd.isnull(run_data['window2_response'].iloc[trial]) and pd.notnull(run_data['window1_response'].iloc[trial]):
                run_data.loc[trial,'calculated_response'] = run_data['window1_response'].iloc[trial]
                run_data.loc[trial,'calculated_RT'] = run_data['window1_RT'].iloc[trial]  
            elif run_data['trial_type'].iloc[trial+1] in control_trials:
                run_data.loc[trial,'calculated_response'] = np.nan
                run_data.loc[trial,'calculated_RT'] = np.nan
            elif pd.isnull(run_data['window2_response'].iloc[trial]) \
                    and (run_data['trial_type'].iloc[trial+1] not in control_trials) and (pd.notnull(run_data['window1_response'].iloc[trial+1])):
                run_data.loc[trial,'calculated_response'] = run_data['window1_response'].iloc[trial+1]
                run_data.loc[trial,'calculated_RT'] = run_data['window1_RT'].iloc[trial+1] + run_data['duration'].iloc[trial]
            elif pd.isnull(run_data['window2_response'].iloc[trial]) \
                    and (run_data['trial_type'].iloc[trial+1] not in control_trials) and (pd.isnull(run_data['window1_response'].iloc[trial+1]))\
                    and (run_data['window2_RT'].iloc[trial+1] < 0.25 and pd.notnull(run_data['window2_response'].iloc[trial+1])):
                run_data.loc[trial,'calculated_response'] = run_data['window2_response'].iloc[trial+1]
                run_data.loc[trial,'calculated_RT'] = run_data['duration'].iloc[trial] + 0.75 + run_data['window2_RT'].iloc[trial+1]
            else:
                run_data.loc[trial,'calculated_response'] = np.nan
                run_data.loc[trial,'calculated_RT'] = np.nan
                                       
    #Determine if calculated response is correct and save to calculated_accuracy column
    for ans in range(len(run_data['correct'].values)):
        if run_data['correct'].iloc[ans] == run_data['calculated_response'].iloc[ans]:
            run_data.loc[ans,'calculated_accuracy'] = 1
        elif run_data['correct'].iloc[ans] != run_data['calculated_response'].iloc[ans]:
            run_data.loc[ans,'calculated_accuracy'] = 0
                 
    #Calculate accuracy, RT, and duration column
    run_data['duration'] = run_data['duration'].round(4)
    #all nan data to 'n/a'
    run_data['calculated_RT'] = run_data['calculated_RT'].round(4)
    
    for col in list(run_data):
        run_data[col] = run_data[col].replace(np.nan, 'n/a')
    
    run_data['window1_response'] = run_data['window1_response_raw']
    
    run_data = run_data.drop(['unique_time'], axis=1)
    
    events = run_data[['onset', 'duration', 'trial_type', 'correct', 'window1_response', 'window1_RT', 'window2_response', 'window2_RT', 'stim_file', 'calculated_accuracy',
                                                'calculated_RT']]
    
    return run_data, events