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

__all__ = ['get_average_oversampling', 'get_oversampling_predicted_sep_r', 'get_oversampling_sep',
           'get_oversampling_predicted_joint_r', 'get_oversampling_predicted_joint']

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

# Cell
def get_average_oversampling(df):
    return df.groupby(['gbe_index']).oversampling.mean().to_frame()

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

def get_oversampling_sep(df):
    # Creating combined variable for ssd and rt
    dfs = []
    # Looping through sessions
    for session in [2,1]:
        # Extracting data for specific session and trial type
        _df = df.query('(session_number==@session)')
        # Predicting scores
        predicted, m = get_oversampling_predicted_sep_r(_df[['participant','oversampling']])
        predicted.columns = ['participant','oversampling_predicted_sep','session']
        # Labeling variables
        predicted['session'] = session
        predicted['gbe_index'] = predicted.participant.astype(str) + '_%03d'%session
        predicted = predicted.set_index('gbe_index')['oversampling_predicted_sep'].to_frame()
# Combining everything into one dataframe
        dfs.append(predicted)
    perc_predicted_sep = pd.concat(dfs)
    return perc_predicted_sep, m


# Cell
def get_oversampling_predicted_joint_r(df):
    R = biuR.wrapper.R()
    p = R("""
    library(lmerTest)
    library(ggeffects)
    # Running the model
    control=lmerControl(optimizer = "bobyqa", optCtrl=list(maxfun=1e6))
    m = lmer(oversampling ~ 1 + (1 | participant / session), data=df, 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_oversampling_predicted_joint(df):
    df['session'] = df.session_number.astype(str) # making session a factor
    # Predicting scores
    predicted, m = get_oversampling_predicted_joint_r(df[['participant','session','oversampling']])
    predicted.columns = ['participant','oversampling_predicted_joint','session']
    #predicted = predicted[('ssrt_predicted_joint','')].to_frame(name="ssrt_predicted_joint").reset_index()
    predicted['gbe_index'] = predicted.participant.astype(str) + predicted.session.apply(lambda x: '_%03d'%int(float(x))).astype(str)
    predicted = predicted.set_index('gbe_index')['oversampling_predicted_joint'].to_frame()
    # Removing sessions that were not in initial dataframe
    predicted = predicted.loc[df.gbe_index.unique()]
    return predicted, m
