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

__all__ = ['DataProvider', 'get_efficiently']

# Cell
from bs4 import BeautifulSoup as bs
import numpy as np
import os
import pandas as pd
from fastcore.foundation import patch

# Cell
class DataProvider():
    def __init__(self, data_folder_path):
        self.data_folder_path = data_folder_path
        self.raw = os.path.join(data_folder_path, 'raw')
        self.external = os.path.join(data_folder_path, 'external')
        self.interim = os.path.join(data_folder_path, 'interim')
        self.processed = os.path.join(data_folder_path, 'processed')
        # Checking if folder paths exist
        assert os.path.isdir(self.external), "External data folder not found."
        assert os.path.isdir(self.raw), "Raw data folder not found."
        assert os.path.isdir(self.interim), "Interim data folder not found."
        assert os.path.isdir(self.processed), "Processed data folder not found."
        self.phonescreening_data_path = os.path.join(self.raw, "phonescreening.csv")
        self.phone_codebook_path = os.path.join(self.external, "phone_codebook.html")
        self.ba_codebook_path = os.path.join(self.external, "ba_codebook.html")
        self.ba_data_path = os.path.join(self.raw, "ba.csv")
        self.b07_participants_path = os.path.join(self.external, "b7_participants.xlsx")


#export
def get_efficiently(func):
    """
    This decorator wraps around functions that get data and handles data storage.
    If the output from the function hasn't been stored yet, it stores it in "[path_to_interim]/[function_name_without_get].parquet"
    If the output from the function has been stored already, it loads the stored file instead of running the function (unless update is specified as True)
    """
    def w(*args, update = False, columns = None, path = None, **kw):
        _self = args[0] # Getting self to grab interim path from DataProvider
        var_name = func.__name__.replace('get_','')
        file_path = os.path.join(_self.interim, "%s.parquet"%var_name)
        if os.path.exists(file_path) and (update == False):
            result =  pd.read_parquet(file_path, columns = columns)
        else:
            print("Preparing %s"%var_name)
            result = func(_self)
            result.to_parquet(file_path)
        return result
    w.__wrapped__ = func # Specifying the wrapped function for inspection
    w.__doc__ = func.__doc__
    w.__name__ = func.__name__
    w.__annotations__ = {'cls':DataProvider, 'as_prop':False} # Adding parameters to make this work with @patch
    return w

# Cell
@patch
def store_interim(self:DataProvider, df, filename):
    path = os.path.join(self.interim,"%s.parquet"%filename)
    df.to_parquet(path)

# Cell
@patch
def load_interim(self:DataProvider, filename):
    return pd.read_parquet(os.path.join(self.interim,"%s.parquet"%filename))

# Cell
@patch
@get_efficiently
def get_phone_codebook(self:DataProvider):
    tables = pd.read_html(open(self.phone_codebook_path,'r').read())
    df = tables[1]
    # Note that str.contains fills NaN values with nan, which can lead to strange results during filtering
    df = df[df.LabelHinweistext.str.contains('Fragebogen:',na=False)==False]
    df = df.set_index('#')
    # Parsing variable name
    df['variable'] = df["Variable / Feldname"].apply(lambda x: x.split(' ')[0])
    # Parsing condition under which variable is displayed
    df['condition'] = df["Variable / Feldname"].apply(lambda x: ' '.join(x.split(' ')[1:]).strip() if len(x.split(' '))>1 else '')
    df['condition'] = df.condition.apply(lambda x: x.replace('Zeige das Feld nur wenn:  ',''))
    # Parsing labels for numerical data
    df['labels'] = np.nan
    labels = tables[2:-1]
    try:
        labels = [dict(zip(l[0],l[1])) for l in labels]
    except:
        display(table)
    searchfor = ["radio","dropdown","yesno","checkbox"]
    with_table = df['Feld Attribute (Feld-Typ, Prüfung, Auswahlen, Verzweigungslogik, Berechnungen, usw.)'].str.contains('|'.join(searchfor))
    df.loc[with_table,'labels'] = labels
    df = df.astype(str)
    return df

# Cell
@patch
def determine_phone_b07(self:DataProvider, df):
    # Some initial fixes
    df.loc[df.center=='d','screen_caller'] = df.loc[df.center=='d','screen_caller'].str.lower().str.strip().replace('leo','leonard visser').replace('sebastian mörcke','sebastian möricke').replace('jessica zimmerman','jessica zimmermann').replace('miriam-sophie petasch','miriam petasch').replace('dorothee','dorothee scheuermann')
    # Cleaning screener list
    dd_screeners = df[(df.center=='d')&(df.screen_caller.isna()==False)].screen_caller.unique()
    def clean_screeners(dd_screeners):
        dd_screeners = [y  for x in dd_screeners for y in x.split('+')]
        dd_screeners = [y  for x in dd_screeners for y in x.split(',')]
        dd_screeners = [y  for x in dd_screeners for y in x.split('und')]
        dd_screeners = [y.replace('(15.02.21)','')  for x in dd_screeners for y in x.split('/')]
        dd_screeners = [y.replace(')','').strip().lower()  for x in dd_screeners for y in x.split('(')]
        dd_screeners = sorted(list(set(dd_screeners)))
        return dd_screeners
    dd_screeners = clean_screeners(dd_screeners)

    b07_screeners = ['ann-kathrin stock','charlotte blum','josephine kirschgens','klara macht','borchardt','marta ledro','miriam petasch','mona hofmann','theo tester']
    s01_screeners = ['esther preuschhof', 'miriam schmitz', 'sebastian möricke', 'jessica zimmermann', 'leonard visser', 'anna-lena lünert', 'anne dörfler', 'dominic reichert', 'maike borchardt', 'dorothee scheuermann', 'paula böhlmann', 'alice']
    known_dd_screeners = list(b07_screeners+s01_screeners)
    dd_screeners = df[(df.center=='d')&(df.screen_caller.isna()==False)].screen_caller.unique()
    # Checking if all Dresden phone screeners are accounted for
    assert df[(df.center=='d')&(df.screen_caller)].screen_caller.str.contains('|'.join(known_dd_screeners)).mean()==1, "Unknown Dresden phone screener: %s"%', '.join(set(clean_screeners(dd_screeners))-set(known_dd_screeners))
    # In general, if a screener from a project was involved, it was screened for that project
    df['screened_for_b07'] = (df.center=='d') & (df.screen_caller.str.contains('|'.join(b07_screeners)))
    df['screened_for_s01'] = (df.center!='d') | (df.screen_caller.str.contains('|'.join(s01_screeners)))

    # We also exclude participants screened for C02 in Berlin
    df.loc[(df.screen_purpose == 4) & (df.center=='b'), 'screened_for_s01'] = False

    # Additionally, we also set it to true if it was specifically set
    df.loc[df.screen_site_dd == 1, 'screened_for_s01'] = True
    df.loc[df.screen_site_dd == 3, 'screened_for_s01'] = True
    df.loc[df.screen_site_dd == 2, 'screened_for_b07'] = True
    df.loc[df.screen_site_dd == 3, 'screened_for_b07'] = True
    return df

# Cell
@patch
def check_participant_id(self:DataProvider,x):
    '''This function checks whether a participant ID is numerical and lower than 20000.'''
    if str(x) == x:
        if x.isnumeric():
            x = float(x)
        else:
            return False
    if x > 20000:
        return False
    return True

# Cell
@patch
def test_check_participant_id(self:DataProvider):
    failed = dp.check_participant_id('test10') == False # Example of a bad participant ID
    passed = dp.check_participant_id('100') == True # Example of a good participant ID
    return failed and passed

# Cell
@patch
def set_dtypes(self:DataProvider, data, codebook):
    '''This function automatically adjust data types of redcap data based on the redcap codebooks'''
    # Parsing type
    codebook['type'] = codebook["Feld Attribute (Feld-Typ, Prüfung, Auswahlen, Verzweigungslogik, Berechnungen, usw.)"].apply(lambda x: x.split(',')[0])
    # Descriptives (not in data)
    desc_columns = list(codebook[codebook.type.str.contains('descriptive')].variable)
    # Datetime
    dt_columns = codebook[(codebook.type.isin(['text (datetime_dmy)','text (date_dmy)']))].variable
    dt_columns = list(set(data.columns).intersection(dt_columns))
    # Numerical
    num_columns = []
    num_columns += list(codebook[codebook.type.str.contains('calc')].variable)
    num_columns += list(codebook[codebook.type.str.contains('checkbox')].variable)
    num_columns += list(codebook[codebook.type.str.contains('radio')].variable)
    num_columns += list(codebook[codebook.type.str.contains('text \(number')].variable)
    num_columns += list(codebook[codebook.type.str.contains('yesno')].variable)
    num_columns += list(codebook[codebook.type.str.contains('dropdown')].variable)
    num_columns += list(codebook[codebook.type.str.contains('slider')].variable)
    num_columns = list(set(data.columns).intersection(num_columns))
    # Text
    text_columns = []
    text_columns += list(codebook[(codebook.type.str.contains('text')) & (~codebook.type.str.contains('date_dmy|datetime_dmy'))].variable)
    text_columns += list(codebook[(codebook.type.str.contains('notes'))].variable)
    text_columns += list(codebook[(codebook.type.str.contains('file'))].variable)
    text_columns = list(set(data.columns).intersection(text_columns))
    assert len(set(num_columns).intersection(set(dt_columns)))==0, set(num_columns).intersection(set(dt_columns))
    assert len(set(text_columns).intersection(set(dt_columns)))==0, set(text_columns).intersection(set(dt_columns))

    for c in num_columns:
        data[c].replace("A 'MySQL server has gone away' error was detected.  It is possible that there was an actual database issue, but it is more likely that REDCap detected this request as a duplicate and killed it.", np.nan, inplace = True)
        data[c] = data[c].astype(float)
    data[text_columns] = data[text_columns].astype(str).replace('nan',np.nan)

    for c in dt_columns:
        data[c] = pd.to_datetime(data[c])
    return data

# Cell
@patch
@get_efficiently
def get_phone_data(self:DataProvider):
    df = pd.read_csv(self.phonescreening_data_path,
                     na_values = ["A 'MySQL server has gone away' error was detected.  It is possible that there was an actual database issue, but it is more likely that REDCap detected this request as a duplicate and killed it."]
                    )
    remove = ['050571', '307493', '345678', '715736', 'Ihloff', 'test',
       'test002', 'test003', 'test004', 'test005', 'test01', 'test02',
       'test03', 'test0722', 'test1', 'test34', 'test999', 'test2020',
       'test20201', 'test345345', 'testt', 'test_10', 'test_11_26',
       'test_neu', 'xx956']
    df = df[~df.participant_id.astype(str).isin(remove)]

    bad_ids = df[~df.participant_id.apply(self.check_participant_id)].participant_id.unique()
    assert len(bad_ids)==0, "Bad participant IDs (should be added to remove): %s"%', '.join(["'%s'"%b for b in bad_ids])
    self.get_phone_codebook()
    df = self.set_dtypes(df, self.get_phone_codebook())
    df['participant_id'] = df.participant_id.astype(int)
    df['center'] = df.screen_site.replace({1:'b',2:'d',3:'m'})
    df['screen_date'] = pd.to_datetime(df['screen_date'], errors = 'coerce')
    #display(df[df.screen_caller.isna()])
    df = self.determine_phone_b07(df)
    return df

# Cell
@patch
@get_efficiently
def get_ba_codebook(self:DataProvider):
    tables = pd.read_html(open(self.ba_codebook_path,"r").read())
    df = tables[1]
    # Note that str.contains fills NaN values with nan, which can lead to strange results during filtering
    df = df[df.LabelHinweistext.str.contains('Fragebogen:',na=False)==False]
    df = df.set_index('#')
    # Parsing variable name
    df['variable'] = df["Variable / Feldname"].apply(lambda x: x.split(' ')[0])
    # Parsing condition under which variable is displayed
    df['condition'] = df["Variable / Feldname"].apply(lambda x: ' '.join(x.split(' ')[1:]).strip() if len(x.split(' '))>1 else '')
    df['condition'] = df.condition.apply(lambda x: x.replace('Zeige das Feld nur wenn:  ',''))
    # Parsing labels for numerical data
    df['labels'] = np.nan
    labels = tables[2:-1]
    try:
        labels = [dict(zip(l[0],l[1])) for l in labels]
    except:
        display(table)
    searchfor = ["radio","dropdown","yesno","checkbox"]
    with_table = df['Feld Attribute (Feld-Typ, Prüfung, Auswahlen, Verzweigungslogik, Berechnungen, usw.)'].str.contains('|'.join(searchfor))
    df.loc[with_table,'labels'] = labels
    df = df.astype(str)
    return df

# Cell
@patch
@get_efficiently
def get_ba_data(self:DataProvider):
    '''This function reads in baseline data from redcap, filters out pilot data, and creates movisens IDs.'''
    df = pd.read_csv(self.ba_data_path)
    df['center'] = df.groupby('participant_id').bx_center.transform(lambda x: x.ffill().bfill())
    df['center'] = df.center.replace({1:'b',2:'d',3:'m'})
    # Creating new movisense IDs (adding center prefix to movisense IDs)
    for old_id in ['bx_movisens','bx_movisens_old','bx_movisens_old_2']:
        new_id = old_id.replace('bx_','').replace('movisens','mov_id')
        df[new_id] = df.groupby('participant_id')[old_id].transform(lambda x: x.ffill().bfill())
        df[new_id] = df.center + df[new_id].astype('str').str.strip('0').str.strip('.').apply(lambda x: x.zfill(3))
        df[new_id].fillna('nan',inplace = True)
        df.loc[df[new_id].str.contains('nan'),new_id] = np.nan
    # Removing test participants
    remove = ['050744', 'hdfghadgfh', 'LindaEngel', 'test', 'Test001', 'Test001a', 'test0011', 'test0012', 'test0013', 'test0014', 'test0015', 'test002', 'test00229', 'test007', 'test01', 'test012', 'test013', 'test1', 'test2', 'test4', 'test12', 'test999', 'test2021', 'test345345', 'testneu', 'testtest', 'test_0720', 'test_10', 'test_GA', 'Test_JH','test0016','891752080', 'pipingTest', 'test0001', 'test00012', 'test0012a', 'test0015a', 'test0017', 'test10', 'test20212', 'testJohn01', 'test_00213', 'test_00233', 'test_00271', 'test_003', 'test_004', 'test_11_26', 'Test_MS']
    df = df[~df.participant_id.astype(str).isin(remove)]
    # Checking participant ids (to find new test participants)
    bad_ids = df[~df.participant_id.apply(self.check_participant_id)].participant_id.unique()
    assert len(bad_ids)==0, "Bad participant IDs (should be added to remove): %s"%', '.join(["'%s'"%b for b in bad_ids])
    # labeling B07 participant
    b07_pps = pd.read_excel(self.b07_participants_path)['Participant ID'].astype(str)
    df['is_b07'] = False
    df.loc[df.participant_id.isin(b07_pps),'is_b07'] = True
    # Setting dtypes based on codebook
    df = self.set_dtypes(df, self.get_ba_codebook())
    # Creating convenience variables
    df['is_female'] = df.screen_gender.replace({1:True,2:False})
    # Filling in missings from baseline
    df['is_female'].fillna(df.bx_sozio_gender.replace({1:False,2:True}), inplace = True)
    df['is_female'] = df.groupby('participant_id')['is_female'].transform(lambda x: x.ffill().bfill())
    return df

# Cell
@patch
def get_duplicate_mov_ids(self:DataProvider):
    '''This function creates a dictionary mapping old to new movisens IDs.'''
    df = self.get_ba_data()
    replace_dict_1 = dict(zip(df.mov_id_old, df.mov_id))
    replace_dict_2 = dict(zip(df.mov_id_old_2, df.mov_id))
    replace_dict = {**replace_dict_1, **replace_dict_2}
    try:
        del replace_dict[np.nan]
    except:
        pass
    return replace_dict