# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/07_gbe.sst.data_provider.ipynb (unless otherwise specified).

__all__ = ['SSTDataProvider']

# Cell
from fastcore.foundation import patch
from ..data_provider import GBEProvider
from ...data_provider import get_efficiently
import numpy as np

# Cell
class SSTDataProvider(GBEProvider):
    '''This class builds upon GBEProvider to get the working memory task data.'''
    def __init__(self, data_folder_path):
        GBEProvider.__init__(self, data_folder_path)

# Cell
@patch
def decode_sst_strings(self:SSTDataProvider, gbe_data):
    df = self.decode_gbe_strings(gbe_data, 'FruitTapGame')
        # Removing left/right distinctions
    df['rt'] = df.lefttime.astype(int) + df.righttime.astype(int)
    df['is_stop'] = (df.stop.astype(int) > 0).astype(float)
    df.loc[df.rt==0,'rt'] = np.nan # Setting 0 RTs to nan
    df['responded'] = (df.rt.isna()==False).astype(float)
    # Calculating SSD
    crw = 650 # ToDo: I'll have to double check this is correct; in Smittenaar it's reported as 500ms, but Ying used 650ms (it's correct as we use the center of response window)
    df['ssd'] = crw - df.gobaddelay.astype(int)
    df.loc[df.is_stop==False,'ssd'] = np.nan
    # Error analysis
    df['omission'] = ((df.is_stop==0) & ((df.rt.isna()) | (df.rt >= 800))).astype(float)
    df['comission'] = ((df.is_stop==1) & (df.rt.isna()==False)).astype(float)
    df['premature'] = (df.rt <= 500).astype(float)
    # Creating convenience variables and restructuring
    df['accuracy'] = df.success.astype(int)
    df = df[[
        'gbe_index',
        'trial_number',
        'anticipation',
        'is_stop','gobaddelay','ssd',
        'responded',
        'rt',
        'accuracy',
        'omission',
        'comission',
        'premature']]
    return df

# Cell
@patch
@get_efficiently
def get_sst_data(self:SSTDataProvider):
    gbe_data = self.get_gbe_data()
    df = self.decode_sst_strings(gbe_data)
    return df