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

__all__ = ['ISTDataProvider']

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

# Cell
class ISTDataProvider(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 get_optimal_turns(self:ISTDataProvider):
    with open(os.path.join(self.external,'opt_string.txt'),'r') as f: #
        optimal_turns = f.read()
    # Getting optimal turns from Hunt's model
    ot = optimal_turns.split('\n')
    ot = [x.split(':') for x in ot][:-1]
    ot = [[list(eval(x[0])),int(x[1])] for x in ot]
    ot = pd.DataFrame(ot)
    ot.columns = ['card_values','optimal_turns']
    ot['deck_id'] = ot.card_values.apply(lambda x: ''.join([str(y-1) for y in x]))
    return ot


# Cell
@patch
def decode_ist_strings(self:ISTDataProvider, gbe_data):
    df = self.decode_gbe_strings(gbe_data, "CardGame")
    back_to_list = lambda x: x.strip('[]').replace('|',',').split(',')
    df['cardlist'] = df.cardlist.apply(back_to_list)
    df['movelist'] = df.movelist.apply(back_to_list)
    df['turntime'] = df.turntime.apply(back_to_list)
    df['deck_id'] = df.cardlist.apply(lambda x: ''.join([str((int(c)-1)%10) for c in x]))
    df['turns'] = df.turntime.apply(lambda x: len(x))
    df['first_card'] = df.deck_id.astype(str).apply(lambda x: int(x[0])+1)
    ot = self.get_optimal_turns()
    df = df.merge(ot[['deck_id','optimal_turns']], on='deck_id',validate='many_to_one').sort_values(by=['gbe_index','trial_number']).reset_index(drop = True)
    df['oversampling'] = df.turns - df.optimal_turns
    return df

# Cell
@patch
@get_efficiently
def get_ist_data(self:ISTDataProvider):
    gbe_data = self.get_gbe_data()
    df = self.decode_ist_strings(gbe_data)
    return df

# Cell
@patch
def filter_old_app_sessions(self:ISTDataProvider, df):
    is_multiplicative = (df.set_index('gbe_index')['trialtype'].astype(int)<2).reset_index().groupby('gbe_index').max()
    sessions_with_old_app = list(is_multiplicative[is_multiplicative.trialtype].index)
    total_sessions = len(df.gbe_index.unique())
    perc_removed = (len(sessions_with_old_app)/total_sessions)*100
    print("In %d (%.2f%%) sessions participants used an old version of the task in some of their sessions.  These sessions  were removed from the dataset."%(len(sessions_with_old_app), perc_removed))
    df = df[df.gbe_index.isin(sessions_with_old_app)==False]
    return df