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

__all__ = ['GBEProvider']

# Cell
import numpy as np
from .data_provider import DataProvider, get_efficiently
from fastcore.foundation import patch
import pandas as pd
from matplotlib import pyplot as plt

# Cell
class GBEProvider(DataProvider):
    '''This class builds upon DataProvider and focusses on GBE data.'''
    def __init__(self, data_folder_path):
        DataProvider.__init__(self, data_folder_path)
        self.GBE_COLUMNS = ['FruitTapGame','WorkingMemoryGame','CardGame','RewardAndHappinessGame']

# Cell
@patch
def get_gbe_from_movisense(self:GBEProvider):
    """
    This function pulls GBE data from movisense and sets cancelled sessions to None.  Here, we only include rows that have at least one non-cancelled GBE game.
    """
    # Getting GBE data
    df = self.get_mov_data().query("Form == 'GreatBrainExperiment'")
    # Setting cancelled sessions to missing
    df[self.GBE_COLUMNS] = df[self.GBE_COLUMNS].replace({'{canceled": true}"':None})
    # Only included rows with at least one non-cancelled GBE game
    df = df[~df[self.GBE_COLUMNS].isna().all(axis=1)]
    df = df.sort_values(['participant','trigger_date'])
    shifted = df.groupby('participant').sampling_day.shift(1)
    df['time_since_last_gbe'] = df['sampling_day'] - shifted
    df['time_since_last_gbe'].fillna(0, inplace = True)
    df['session_number'] = df.groupby('participant').cumcount() + 1
    # Creating the trigger type variable
    replace_dict = {
    'Button Pressed Spiele starten und Initialfragen (~40min)':'initial',
    'Participant Defined Time Trigger':'ema',
    'Button Pressed Bedarfstart GBE (Ausnahme)':'optional',
    'Button Pressed Bedarfsstart GBE':'optional'}
    df['trigger_type'] = df.Trigger.apply(lambda x: x.replace('Pressed:','Pressed').split(':')[0]).replace(replace_dict)
    df['gbe_index'] = df.participant + '_' + df.session_number.apply(lambda x: '%03d'%int(x))
    df = df.set_index('gbe_index')
    return df

# Cell
@patch
def define_initial_dataset(self:GBEProvider, df):
    initial_pps = df.sort_values(['starting_date','trigger_date']).groupby('participant').first().iloc[:300].index
    df['is_initial'] = df.participant.isin(initial_pps)
    df = df.sort_index()
    return df

# Cell
@patch
def get_gbe_data(self:GBEProvider):
    df = self.get_gbe_from_movisense()
    df = self.define_initial_dataset(df)
    return df

# Cell
@patch
def define_baseline_sessions(self:GBEProvider, df):
    first_two_sessions = df.session_number <= 2
    same_day = df.time_since_last_gbe == 0
    initial_or_optional = df.trigger_type != 'ema'
    df['is_baseline'] = first_two_sessions & same_day & initial_or_optional
    return df