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

__all__ = ['get_max_levels', 'get_perc_correct', 'add_implied_data', 'get_implied_perc_correct',
           'get_perc_correct_predicted_sep_r', 'get_perc_correct_predicted_sep',
           'get_perc_correct_predicted_sep_trial_r', 'get_perc_correct_predicted_sep_trial']

# Cell
from .data_provider import WMDataProvider
import pandas as pd
import numpy as np
from scipy import stats
import biuR.wrapper

# Cell
def get_max_levels(df):
    '''This function calculates max levels from wm data.'''
    # Selecting successful trials only.
    successful_trials = df.query('success==1')
    # Grouping by participant, session, and condition and getting the maximum level.
    max_levels = successful_trials.groupby(['gbe_index','trial_type'])['level'].max().unstack()
    # Labeling variables
    max_levels = max_levels[['no_distractor_1', 'no_distractor_2', 'encoding_distractor', 'delayed_distractor']]
    max_levels = max_levels.add_prefix("max_level_")
    return max_levels

# Cell
def get_perc_correct(df):
    '''This function calculates percentage correct from wm data.'''
    max_levels = df.groupby(['gbe_index','trial_type'])['success'].mean().unstack()
    max_levels = max_levels[['no_distractor_1', 'no_distractor_2', 'encoding_distractor', 'delayed_distractor']]
    max_levels = max_levels.add_prefix("perc_correct_")
    return max_levels

# Cell
def add_implied_data(df):
    # Checking whether participants needed one or two trials to complete a session
    df['duplication'] = df.groupby(['participant','session_number','trial_type','level']).cumcount()
    _df = df.set_index(['participant','session_number','trial_type','level','duplication']).sort_index()[['success']]
    # Imputing data
    levels = list(range(3,11))
    new_index = pd.MultiIndex.from_product([_df.reset_index().participant.unique(), _df.reset_index().session_number.unique(),_df.reset_index().trial_type.unique(), levels, [0,1]],names = ['participant','session_number','trial_type','level','duplication'])
    _df = _df.reindex(new_index)
    _df['implied_success'] = _df.groupby(level=['participant','session_number','trial_type','level']).transform(np.max).fillna(0)
    _df['implied_success'] = _df.success.fillna(_df.implied_success)
    _df = _df.reset_index()
    #_df['implied_success'] = _df.implied_success# * _df.level
    _df = _df.merge(df.groupby(['participant','session_number']).gbe_index.first().reset_index(), on = ['participant','session_number'], how = 'outer' )
    _df = _df[_df.gbe_index.isnull()==False]
    #implied_success = _df.groupby(['gbe_index']).implied_success.mean()
    return _df#, implied_success

# Cell
def get_implied_perc_correct(df):
    df = add_implied_data(df)
    '''This function calculates percentage correct from wm data.'''
    implied_perc_correct = df.groupby(['gbe_index','trial_type'])['implied_success'].mean().unstack()
    implied_perc_correct = implied_perc_correct[['no_distractor_1', 'no_distractor_2', 'encoding_distractor', 'delayed_distractor']]
    implied_perc_correct = implied_perc_correct.add_prefix("impl_perc_correct_")
    return implied_perc_correct

# Cell
#def get_perc_correct_predicted_sep_r(df):
#    %R -i df
#    %R library(lmerTest)
#    %R library(ggeffects)
#    # Running the model
#    %R control=glmerControl(optimizer = "bobyqa", optCtrl=list(maxfun=1e6))
#    %R m = glmer(implied_success ~ 1 + (1 | participant), data=df, family=binomial, control = control, na.action = na.exclude)
#    # Extracting predicted values
#    %R p = ggpredict(m, terms=c("participant"), type="re",ci.lvl = NA)
#    %R -o p
#    return p

def get_perc_correct_predicted_sep_r(df):
    R = biuR.wrapper.R()
    p = R("""
    library(lmerTest)
    library(ggeffects)
    # Running the model
    control=glmerControl(optimizer = "bobyqa", optCtrl=list(maxfun=1e6))
    m = glmer(implied_success ~ 1 + (1 | participant), data=df, family=binomial, control = control, na.action = na.exclude)
    # Extracting predicted values
    ggpredict(m, terms=c("participant"), type="re",ci.lvl = NA)
    """,push=dict(df=df))
    return p

def get_perc_correct_predicted_sep(df):
    df = add_implied_data(df)
    df = df[['implied_success','participant','session_number','trial_type']]

    dfs = []
    # Looping through trial types
    for trial_type in ['no_distractor_2','no_distractor_1','encoding_distractor','delayed_distractor']:
        session_dfs = []
        # Looping through sessions
        for session in [2,1]:
            # Extracting data for specific session and trial type
            _df = df.query('(session_number==@session) and (trial_type==@trial_type)')[['implied_success','participant']]
            # Predicting scores
            predicted = get_perc_correct_predicted_sep_r(_df)
            # Labeling variables
            predicted.columns = ['participant','perc_predicted_sep_%s'%trial_type,'session']
            predicted['session'] = session
            predicted['gbe_index'] = predicted.participant.astype(str) + '_%03d'%session
            predicted = predicted.set_index('gbe_index')['perc_predicted_sep_%s'%trial_type].to_frame()
    # Combining everything into one dataframe
            session_dfs.append(predicted)
        dfs.append(pd.concat(session_dfs))
    perc_predicted_sep = pd.concat(dfs, axis = 1)
    return perc_predicted_sep

# Cell
def get_perc_correct_predicted_sep_trial_r(df):
    R = biuR.wrapper.R()
    p = R("""
    library(lmerTest)
    library(ggeffects)
    # Running the model
    control=glmerControl(optimizer = "bobyqa", optCtrl=list(maxfun=1e6))
    m = glmer(implied_success ~ 1 + (1 | participant/session), data=df, family=binomial, control = control, na.action = na.exclude)
    # Extracting predicted values
    ggpredict(m, terms=c("participant","session"), type="re",ci.lvl = NA)
    """,push=dict(df=df))

    m = R("""m""")
    return p, m

def get_perc_correct_predicted_sep_trial(df):
    df = add_implied_data(df)
    dfs = []
    ms = {}
    # Looping through trial types
    for trial_type in ['no_distractor_2','no_distractor_1','encoding_distractor','delayed_distractor']:
        # Extracting data for specific trial type
        _df = df.query('(trial_type==@trial_type)')[['implied_success','participant','session_number']]
        _df['session'] = _df.session_number.astype(str) # making session a factor
        # Predicting scores
        predicted, m = get_perc_correct_predicted_sep_trial_r(_df)
        # Labeling variables
        predicted.columns = ['participant','perc_predicted_sep_trial_%s'%trial_type,'session']
        #predicted['session'] = session
        predicted['gbe_index'] = predicted.participant.astype(str) + predicted.session.apply(lambda x: '_%03d'%int(float(x))).astype(str)
        predicted = predicted.set_index('gbe_index')['perc_predicted_sep_trial_%s'%trial_type].to_frame()
    # Combining everything into one dataframe
        dfs.append(predicted)
        ms[trial_type] = m
    perc_predicted_sep_trial = pd.concat(dfs, axis = 1)
    # Removing sessions that were not in initial dataframe
    perc_predicted_sep_trial = perc_predicted_sep_trial.loc[df.gbe_index.unique()]
    return perc_predicted_sep_trial, ms
