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

__all__ = ['GBEProvider', 'determine_ema_session_numbers', 'check_repeating_sessions', 'increase_session_numbers',
           'filter_repeating_sessions']

# 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
import re

# 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',
    'Repeated':'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.groupby('participant').first().sort_values(['starting_date','trigger_date']).iloc[:300].index
    df['is_initial'] = df.participant.isin(initial_pps)
    df = df.sort_index()
    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'
    before_ema = df.sampling_day < 22 # Ema starts on sampling day 22, so we only consider baseline sessions that happened before that
    df['is_baseline'] = first_two_sessions & same_day & initial_or_optional & before_ema
    return df

# Cell
@patch
def find_gbe_game_starting_time(self:GBEProvider, s):
    '''As baseline GBEs have the same trigger date, we first have to parse the time from the GBE data.'''
    nan_time = pd.to_datetime(np.nan)
    try:
        gbe_columns = s.split(',')

    except:
        return nan_time
    if 'timestarted' not in gbe_columns:
        return nan_time
    time_started_index = gbe_columns.index('timestarted')
    time_started = pd.to_datetime(s.split('\n')[1].split(',')[time_started_index][2:-2])
    return time_started



@patch
def check_time_between_sessions(self:GBEProvider, df):
    time_1 = pd.to_datetime(np.nan)
    time_2 = pd.to_datetime(np.nan)
    if len(df)==2:
        time_1 = self.find_gbe_game_starting_time(df.iloc[0])
        time_2 = self.find_gbe_game_starting_time(df.iloc[1])
    try:
        return (time_2 - time_1).seconds/60
    except:
        return np.nan

# Cell
def determine_ema_session_numbers(_df):
    '''Ema sessions start the earliest after sampling day 22.  The following year is split into 28-day intervals, each corresponding to one EMA session.'''
    _df['ema_session'] = pd.cut(_df.sampling_day, np.arange(21, 392, 28), labels = [x for x in range(1,14)]).astype(float)
    return _df

def check_repeating_sessions(_df):
    '''Sometimes EMA sessions end up happening in the same 28 days interval.  In these cases, participants therefore have repeating sessions.  If a repeating session is one the same day, it should certainly be removed.  However, if a repeating session happens far enough from other sessions and the subsequent session is missing, we can assume that this session is actually a real session.  Although, the nominal distance between sessions is 28 days, participants can postpone sessions, thus reducing the time to the next session.  Specifying at least 14 one days between sessions, therefore seems reasonable.'''
    _df['ema_session_repeating'] = ((_df.ema_session.astype(float) - _df.ema_session.astype(float).shift(1))==0)
    _df['total_ema_sessions_completed'] = _df.groupby('participant').ema_session.transform(lambda x:len(x.dropna().unique()))
    _df['time_since_last_session'] = ((_df.sampling_day - _df.groupby('participant').sampling_day.shift(1)))
    _df['time_to_next_session'] = ((_df.sampling_day - _df.groupby('participant').sampling_day.shift(-1))*-1)
    # Deciding wether to increase session number
    sessions_missing = _df['total_ema_sessions_completed'] < 13
    window = 14
    enough_distance = (_df.time_since_last_session>window) & (_df.time_to_next_session>window)
    _df['increase_session_number'] = (_df.ema_session_repeating & enough_distance & sessions_missing).astype(int)
    return _df

def increase_session_numbers(_df):
    _df['ema_session'] = _df.ema_session + _df.increase_session_number
    return(_df)

def filter_repeating_sessions(_df):
    _df = _df.query('ema_session_repeating==False')
    _df = _df.query('is_baseline or (ema_session.isna() == False)')
    return _df

@patch
def fix_repeating_ema_sessions(self:GBEProvider, _df):
    _df = determine_ema_session_numbers(_df)
    old_df = _df
    for i in range(13):
        _df = check_repeating_sessions(_df)
        _df = increase_session_numbers(_df)
    old_sessions = _df.index

    _df = filter_repeating_sessions(_df)
    new_sessions = _df.index
    removed_sessions = list((set(old_sessions)-set(new_sessions)))
    #return old_df.loc[removed_sessions]
    removed_participants = list(set([s.split('_')[0] for s in removed_sessions]))
    _df = filter_repeating_sessions(_df)
    _df = check_repeating_sessions(_df)
    print("%d participants had unplanned sessions.  %d sessions (%.2f%%) were removed."%(len(removed_participants), len(removed_sessions), len(old_df)/len(removed_sessions)))
    _df['session_number_ex_protocol'] = _df.session_number
    _df['session_number'] = (_df.ema_session+2).fillna(_df.session_number)
    return _df

# Cell
@patch
@get_efficiently
def get_gbe_data(self:GBEProvider):
    df = self.get_gbe_from_movisense()
    df = self.define_initial_dataset(df)
    print(df.is_initial.isna().mean())
    df = self.define_baseline_sessions(df)
    print(df.is_initial.isna().mean())
    df = self.fix_repeating_ema_sessions(df)
    print(df.is_initial.isna().mean())
    df = df.reset_index()
    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 get_initial_baseline(self:GBEProvider):
    df = self.get_gbe_data()
    return df.query("is_initial and is_baseline")

@patch
def get_replication_baseline(self:GBEProvider):
    df = self.get_gbe_data()
    return df.query("(is_baseline and (not initial)")

# Cell
@patch
def decode_gbe_string(self:GBEProvider, s):
    """This function function turns one gbe output string into a dataframe"""
    def replace(g):
        return g.group(0).replace(',', '|')
    s = re.sub(r'\[.*?\]', replace, s) # The comma separated string can contain comma separated list items
    columns, df = s.replace('","',';').replace('"','').split('\n')
    df = pd.DataFrame([column.split(',') for column in df.split(';')][:-1]).transpose().ffill().iloc[:-1]
    df.columns = [c.replace('tr_','') for c in columns.split(',')[:-1]]
    def to_datetime(x):
        return pd.to_datetime(x[1:-1])
    if 'timestarted' in df.columns:
        df['timestarted'] = df.timestarted.apply(to_datetime)
        df['timesubmitted'] = df.timesubmitted.apply(to_datetime)
        df.drop(columns = 'appversion', inplace = True)
    else:
        df['timestarted'] = pd.to_datetime(np.nan)
        df['timesubmitted'] = pd.to_datetime(np.nan)
    return df

# Cell
@patch
def decode_gbe_strings(self:GBEProvider, df, column):
    '''This function turns all gbe output strings in a column into dataframes and concatenates them.'''
    df = df[~df[column].isna()] # Selecting non nan data
    #gbe_data = pd.concat(df.set_index(['participant','session_number'])[column].apply(self.decode_gbe_string).values, keys = df.index)
    #gbe_data = pd.concat(df[column].apply(self.decode_gbe_string).values)
    df = pd.concat(df[column].apply(self.decode_gbe_string).values, keys = df.index)
    df.index.rename('trial_number',level = 1, inplace = True)
    df = df.reset_index()
    df['trial_number'] = df.trial_number + 1
    return df