# AUTOGENERATED! DO NOT EDIT! File to edit: 00_core.ipynb (unless otherwise specified).

__all__ = ['load_public_data', 'prepare_baseline_and_intervention_usable_data', 'get_adherent_column',
           'most_active_user', 'convert_loggings', 'get_certain_types', 'breakfast_analysis_summary',
           'breakfast_analysis_variability', 'dinner_analysis_summary', 'dinner_analysis_variability']

# Cell
import pandas as pd
import numpy as np
from scipy import stats
import seaborn as sns
import os
import matplotlib.pyplot as plt
import pickle
from datetime import date
from datetime import datetime
from collections import defaultdict

# Cell
def load_public_data(in_path, export = False, out_path = 'data/output/public.pickle'):
    """
    Load original public data and output processed data in pickle format.\n

    Process include:\n
    1. Dropping 'foodimage_file_name' column.\n
    2. Handling the format of time by generating a new column, 'original_logtime_notz'\n
    3. Generating the date column, 'date'\n
    4. Converting time into float number into a new column, 'local_time'\n
    5. Converting time in the 'local_time' column so that day starts at 4 am.\n
    6. Converting time to a format of HH:MM:SS, 'time'\n
    7. Generating the column 'week_from_start' that contains the week number that the participants input the food item.\n
    8. Generating 'year' column based on the input data.\n
    9. Outputing the data into a pickle format file.\n

    \n
    @param in_path : input path\n
    @param out_path: output path\n
    @param export : whether to save the processed dataframe locally\n
    @return: processed dataframe
    """
    public_all = pd.read_csv(in_path).drop(columns = ['foodimage_file_name'])

    def handle_public_time(s):
        tmp_s = s.replace('p.m.', '').replace('a.m.', '')
        try:
            return pd.to_datetime(' '.join(tmp_s.split()[:2]) )
        except:
            try:
                if int(tmp_s.split()[1][:2]) > 12:
                    tmp_s = s.replace('p.m.', '').replace('a.m.', '').replace('PM', '').replace('pm', '')
                return pd.to_datetime(' '.join(tmp_s.split()[:2]) )
            except:
                return np.nan

    original_logtime_notz_lst = []
    for t in (public_all.original_logtime.values):
        original_logtime_notz_lst.append(handle_public_time(t))
    public_all['original_logtime_notz'] = original_logtime_notz_lst

    public_all = public_all.dropna().reset_index(drop = True)

    def find_date(d):
        if d.hour < 4:
            return d.date() - pd.Timedelta('1 day')
        else:
            return d.date()
    public_all['date'] = public_all['original_logtime_notz'].apply(find_date)


    # Handle the time - Time in floating point format
    public_all['local_time'] = public_all.original_logtime_notz.apply(lambda x: pd.Timedelta(x.time().isoformat()).total_seconds() /3600.).values
    day_begins_at = 4
    public_all.loc[(public_all['local_time'] < day_begins_at), 'local_time'] = 24.0 + public_all.loc[(public_all['local_time'] < day_begins_at), 'local_time']

    # Handle the time - Time in Datetime object format
    public_all['time'] = pd.DatetimeIndex(public_all.original_logtime_notz).time

    # Handle week from start
    public_start_time_dic = dict(public_all.groupby('unique_code').agg(np.min)['date'])
    def count_week_public(s):
        return (s.date - public_start_time_dic[s.unique_code]).days // 7 + 1
    public_all['week_from_start'] = public_all.apply(count_week_public, axis = 1)

    public_all['year'] = public_all.date.apply(lambda d: d.year)

    if export == True:
        public_all_pickle_file = open(out_path, 'wb')
        pickle.dump(public_all, public_all_pickle_file)
        print('data is saved at {}'.format(out_path))
        public_all_pickle_file.close()

    return public_all

# Cell
def prepare_baseline_and_intervention_usable_data(in_path,
                                                  export = False,
                                                  baseline_expanded_out_path='data/output/public_basline_usable_expanded.pickle',
                                                 intervention_out_path = 'data/output/public_intervention_usable.pickle'):
    """
    Filter and create baseline_expanded and intervention groups based on in_path pickle file.\n
    @param in_path : input path, file in pickle, csv or pandas dataframe format\n
    @param out_path: output path\n
    @param export : whether to save the processed dataframes locally\n
    @return: baseline expanded and intervention dataframes in a list format where index 0 is the baseline dataframe and 1 is the intervention dataframe.
    """

    if isinstance(in_path, str):
        if in_path.split('.')[-1] == 'pickle':
            # load data
            public_all_pickle_file = open(in_path, 'rb')
            public_all = pickle.load(public_all_pickle_file)
            print('read the pickle file successfully.')
        if in_path.split('.')[-1] == 'csv':
            public_all = pd.read_csv(in_path)
            print('read the csv file successfully.')
    else:
        public_all = in_path

    # create baseline data
    df_public_baseline = public_all.query('week_from_start <= 2')
    df_public_baseline_usable, public_baseline_usable_id_set = \
    filtering_usable_data(df_public_baseline, num_items = 40, num_days = 12)

    # create intervention data
    df_public_intervention = public_all.query('week_from_start in [13, 14]')
    df_public_intervention_usable, public_intervention_usable_id_set = \
    filtering_usable_data(df_public_intervention, num_items = 20, num_days = 8)

    # create df that contains both baseline and intervention id_set that contains data for the first two weeks
    expanded_baseline_usable_id_set = set(list(public_baseline_usable_id_set) + list(public_intervention_usable_id_set))
    df_public_basline_usable_expanded = public_all.loc[public_all.apply(lambda s: s.week_from_start <= 2 \
                                                    and s.unique_code in expanded_baseline_usable_id_set, axis = 1)]

    if export == True:

        # save baseline_expanded and intervention to pickle_file
        public_pickle_file = open(baseline_expanded_out_path, 'wb')
        pickle.dump(df_public_basline_usable_expanded, public_pickle_file)
        public_pickle_file.close()
        print('baseline_expanded data is saved at {}'.format(baseline_expanded_out_path))

        public_pickle_file = open(intervention_out_path, 'wb')
        pickle.dump(df_public_intervention_usable, public_pickle_file)
        public_pickle_file.close()
        print('intervention data is saved at {}'.format(intervention_out_path))

    return [df_public_basline_usable_expanded, df_public_intervention_usable]

# Cell
def get_adherent_column(in_path, out_path=None, export = False):
    """
    @param in_path: input path, file in pickle, csv or pandas dataframe format\n
    @param out_path: path for out_put file in csv format if export==True.\n
    @param export: whether to save processed df locally.\n
    return True for a logging if the there are more than 2 loggings in one day w/ more than 4hrs apart from each other.
    """
    def adherent(s):
        if len(s.values) >= 2 and (max(s.values) - min(s.values)) >= 4:
            return True
        else:
            return False

    if isinstance(in_path, str):
        if in_path.split('.')[-1] == 'pickle':
            # load data
            public_all_pickle_file = open(in_path, 'rb')
            df = pickle.load(public_all_pickle_file)
            print('read the pickle file successfully.')
        if in_path.split('.')[-1] == 'csv':
            df = pd.read_csv(in_path)
            print('read the csv file successfully.')
    else:
        df = in_path

    adherent_dict = dict(df.groupby(['unique_code', 'date'])['local_time'].agg(adherent))
    df['adherence'] = df.apply(lambda x: adherent_dict[(x.unique_code, x.date)], axis = 1)

    if export == True:
        df.to_csv(out_path)

    return df

# Cell
def most_active_user(in_path, n, export = False, user_day_counts_path = 'data/output/public_top_users_day_counts.csv'):
    """
    @param in_path : input path, file in pickle, csv or pandas dataframe format\n
    @param n: number of users with most number of logging days\n
    @param user_data_counts_path: output path for top_users_day_counts in csv format\n
    @param export : whether to save the processed dataframe locally\n
    @return: top_users_day_counts dataframe\n

    This function returns a dataframe in csv format that finds top n users with the most number of days that they logged.
    """
    if isinstance(in_path, str):
        if in_path.split('.')[-1] == 'pickle':
            # load data
            public_all_pickle_file = open(in_path, 'rb')
            public_all = pickle.load(public_all_pickle_file)
            print('read the pickle file successfully.')
        if in_path.split('.')[-1] == 'csv':
            public_all = pd.read_csv(in_path)
            print('read the csv file successfully.')
    else:
        public_all = in_path

    # filter the dataframe so it only contains food and beverage food type, then get the top n users who input the
    # most loggings in descending order
    top_users = public_all.query('food_type in ["f", "b"]')[['ID', 'unique_code']]\
    .groupby('unique_code').agg('count').sort_values(by = 'ID', ascending = False).index[:n]
    top_users = set(list(top_users))
    public_top_users = public_all.query('food_type in ["f", "b"]')\
    .loc[public_all.unique_code.apply(lambda x: x in top_users)].reset_index(drop = True)


    public_top_users = get_adherent_column(public_top_users)

    public_top_users_day_counts = pd.DataFrame(public_top_users.query('adherence == True')\
                            [['date', 'unique_code']].groupby('unique_code')['date'].nunique())\
                            .sort_values(by = 'date', ascending = False).rename(columns = {'date': 'day_count'})

    if export == True:
        public_top_users_day_counts.to_csv(user_day_counts_path)
        print('user day counts data is saved at {}'.format(user_day_counts_path))

    return public_top_users_day_counts


# Cell
def convert_loggings(in_path, export = False, parsed_food_path='data/output/public_all_parsed.csv', ascending = False):
    """
    @param in_path : input path, file in pickle, csv or panda dataframe format\n
    @param parsed_food_path : output path for cleaned food loggings in csv format \n
    @param ascending: True to sort the item from lowest to highest frequency and vice versa\n
    @param export: whether to save the converted dataframe locally\n
    @return: frequency count of food loggings in the descending order by default\n

    This function takes in a pickle file, convert all the loggings into individual items of food and count the frequency of all food loggings.
    """
    if isinstance(in_path, str):
        if in_path.split('.')[-1] == 'pickle':
            # load data
            public_all_pickle_file = open(in_path, 'rb')
            public_all = pickle.load(public_all_pickle_file)
            print('read the pickle file successfully.')
        if in_path.split('.')[-1] == 'csv':
            public_all = pd.read_csv(in_path)
            print('read the csv file successfully.')
    else:
        public_all = in_path

    # initialize food parser instance
    fp = FoodParser()
    fp.initialization()

    # parse food
    parsed = [fp.parse_food(i, return_sentence_tag = True) for i in public_all.desc_text.values]

    print('All food is parsed!')

    public_all_parsed = pd.DataFrame({
    'ID': public_all.unique_code,
    'food_type': public_all.food_type,
    'desc_text': public_all.desc_text,
    'cleaned': parsed
    })

    public_all_parsed['cleaned'] = public_all_parsed['cleaned'].apply(lambda x: x[0])

    if export == True:
        public_all_parsed.to_csv(parsed_food_path)
        print('cleaned food loggings is saved at {}'.format(parsed_food_path))

    # count food
    all_count_dict = defaultdict(lambda : 0)
    for food in public_all_parsed.cleaned.values:
        if len(food) > 0:
            for f in food:
                all_count_dict[f] += 1

    print('All food is counted!')

    all_count_dict = dict(all_count_dict)
    public_freq_count = pd.DataFrame({
        'food': list(all_count_dict.keys()),
        'count': [all_count_dict[k] for k in list(all_count_dict.keys())]
    }).sort_values(by = 'count', ascending = ascending).reset_index(drop = True)

    return public_freq_count

# Cell
def get_certain_types(in_path, types, out_path=None, export = False):
    """
    @param in_path : input path, file in pickle, csv or panda dataframe format\n
    @param types : expected types of the loggings for filtering, in format of list. Available types:  \n
    1. 'w' : water \n
    2. 'b' : beverage \n
    3. 'f' : food \n
    4. 'm' : medicine \n
    @param export: whether to save the dataframe locally\n
    @param out_path: path of the saved file if export==True, csv format\n
    @return: filtered dataframe with expected type/types with four columns: 'unique_code', 'desc_text', 'date', 'local_time'
    """
    if isinstance(in_path, str):
        if in_path.split('.')[-1] == 'pickle':
            df = pickle.load(open(in_path,'rb'))
            print('read the pickle file successfully.')
        if in_path.split('.')[-1] == 'csv':
            df = pd.read_csv(in_path)
            print('read the csv file successfully.')
    else:
        df = in_path

    if len(types) == 0:
        return df[['unique_code','food_type', 'desc_text', 'date', 'local_time']]

    if len(types) == 1:
        if types[0] not in ['w', 'b', 'f', 'm']:
            raise Exception("not a valid logging type")
        filtered = df[df['food_type']==types[0]][['unique_code','food_type', 'desc_text', 'date', 'local_time']].reset_index(drop = True)
    else:
        filtered = df[df['food_type'].isin(types)][['unique_code','food_type', 'desc_text', 'date', 'local_time']].reset_index(drop = True)

    if export == True:
        filtered.to_csv(out_path)
        print('filtered dataframe is saved at {}.'.format(out_path))

    return filtered


# Cell
def breakfast_analysis_summary(in_path, out_path=None, export = False):
    """
    @param in_path: input path, file in pickle, csv or panda dataframe format\n
    @param export: whether to save the dataframe locally\n
    @param out_path: path of the saved file if export==True, csv format\n
    @return: 5%,10%,25%,50%,75%,90%,95% quantile of breakfast time for all subjects from the in_path file
    """

    if isinstance(in_path, str):
        if in_path.split('.')[-1] == 'pickle':
            df = pickle.load(open(in_path,'rb'))
            print('read the pickle file successfully.')
        if in_path.split('.')[-1] == 'csv':
            df = pd.read_csv(in_path)
            print('read the csv file successfully.')
    else:
        df = in_path

    # leave only the adherent loggings
    df = get_adherent_column(df)
    df = df[df['adherence']==True]

    breakfast_series = df.groupby(['unique_code', 'date'])['local_time'].min().groupby('unique_code').quantile([0.05, 0.10, 0.25, 0.5, 0.75, 0.90, 0.95])
    breakfast_df = pd.DataFrame(breakfast_series)
    all_rows = []
    for index in breakfast_df.index:
        tmp_dict = dict(breakfast_series[index[0]])
        tmp_dict['id'] = index[0]
        all_rows.append(tmp_dict)
    breakfast_summary_df = pd.DataFrame(all_rows, columns = ['id', 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95])\
        .rename(columns = {0.05: '5%', 0.1: '10%', 0.25: '25%', 0.5: '50%', 0.75: '75%', 0.9: '90%', 0.95: '95%'})\
        .drop_duplicates().reset_index(drop = True)

    if export == True:
        breakfast_summary_df.to_csv(out_path)

    return breakfast_summary_df

# Cell
def breakfast_analysis_variability(in_path, out_path=None, export = False):
    """
    @param in_path: input path, file in pickle, csv or panda dataframe format\n
    @param export: whether to save the dataframe locally\n
    @param out_path: path of the saved file if export==True, csv format\n
    @return: returns a dataframe that contains 5%,10%,25%,50%,75%,90%,95% quantile of dinner time minus 50% time for all subjects from the in_path file.\n
    Also returns a histgram that represent the 90%-10% interval for all subjects.
    """

    if isinstance(in_path, str):
        if in_path.split('.')[-1] == 'pickle':
            df = pickle.load(open(in_path,'rb'))
            print('read the pickle file successfully.')
        if in_path.split('.')[-1] == 'csv':
            df = pd.read_csv(in_path)
            print('read the csv file successfully.')
    else:
        df = in_path

    # leave only the adherent loggings
    df = get_adherent_column(df)
    df = df[df['adherence']==True]

    breakfast_series = df.groupby(['unique_code', 'date'])['local_time'].min().groupby('unique_code').quantile([0.05, 0.10, 0.25, 0.5, 0.75, 0.90, 0.95])
    breakfast_df = pd.DataFrame(breakfast_series)
    all_rows = []
    for index in breakfast_df.index:
        tmp_dict = dict(breakfast_series[index[0]])
        tmp_dict['id'] = index[0]
        all_rows.append(tmp_dict)
    breakfast_summary_df = pd.DataFrame(all_rows, columns = ['id', 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95])\
        .rename(columns = {0.05: '5%', 0.1: '10%', 0.25: '25%', 0.5: '50%', 0.75: '75%', 0.9: '90%', 0.95: '95%'})\
        .drop_duplicates().reset_index(drop = True)
    breakfast_variability_df = breakfast_summary_df.copy()

    for col in breakfast_variability_df.columns:
        if col == 'id' or col == '50%':
            continue
        breakfast_variability_df[col] = breakfast_variability_df[col] - breakfast_variability_df['50%']
    breakfast_variability_df['50%'] = breakfast_variability_df['50%'] - breakfast_variability_df['50%']

    fig, ax = plt.subplots(1, 1, figsize = (10, 10), dpi=80)
    sns_plot = sns.distplot( breakfast_variability_df['90%'] - breakfast_variability_df['10%'] )
    ax.set(xlabel='Variation Distribution for Breakfast (90% - 10%)', ylabel='Kernel Density Estimation')

    if export == True:
        breakfast_variability_df.to_csv(out_path)

    return breakfast_variability_df

# Cell
def dinner_analysis_summary(in_path, out_path=None, export = False):
    """
    @param in_path: input path, file in pickle, csv or panda dataframe format\n
    @param export: whether to save the dataframe locally\n
    @param out_path: path of the saved file if export==True, csv format\n
    @return: 5%,10%,25%,50%,75%,90%,95% quantile of dinner time for all subjects from the in_path file
    """

    if isinstance(in_path, str):
        if in_path.split('.')[-1] == 'pickle':
            df = pickle.load(open(in_path,'rb'))
            print('read the pickle file successfully.')
        if in_path.split('.')[-1] == 'csv':
            df = pd.read_csv(in_path)
            print('read the csv file successfully.')
    else:
        df = in_path

    # leave only the adherent loggings
    df = get_adherent_column(df)
    df = df[df['adherence']==True]

    dinner_series = df.groupby(['unique_code', 'date'])['local_time'].max().groupby('unique_code').quantile([0.05, 0.10, 0.25, 0.5, 0.75, 0.90, 0.95])
    dinner_df = pd.DataFrame(dinner_series)
    all_rows = []
    for index in dinner_df.index:
        tmp_dict = dict(dinner_series[index[0]])
        tmp_dict['id'] = index[0]
        all_rows.append(tmp_dict)
    dinner_summary_df = pd.DataFrame(all_rows, columns = ['id', 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95])\
        .rename(columns = {0.05: '5%', 0.1: '10%', 0.25: '25%', 0.5: '50%', 0.75: '75%', 0.9: '90%', 0.95: '95%'})\
        .drop_duplicates().reset_index(drop = True)

    if export == True:
        dinner_summary_df.to_csv(out_path)

    return dinner_summary_df

# Cell
def dinner_analysis_variability(in_path, out_path=None, export = False):
    """
    @param in_path: input path, file in pickle, csv or panda dataframe format\n
    @param export: whether to save the dataframe locally\n
    @param out_path: path of the saved file if export==True, csv format\n
    @return: returns a dataframe that contains 5%,10%,25%,50%,75%,90%,95% quantile of dinner time minus 50% time for all subjects from the in_path file. Also returns a histgram that represent the 90%-10% interval for all subjects.
    """

    if isinstance(in_path, str):
        if in_path.split('.')[-1] == 'pickle':
            df = pickle.load(open(in_path,'rb'))
            print('read the pickle file successfully.')
        if in_path.split('.')[-1] == 'csv':
            df = pd.read_csv(in_path)
            print('read the csv file successfully.')
    else:
        df = in_path

    # leave only the adherent loggings
    df = get_adherent_column(df)
    df = df[df['adherence']==True]

    dinner_series = df.groupby(['unique_code', 'date'])['local_time'].max().groupby('unique_code').quantile([0.05, 0.10, 0.25, 0.5, 0.75, 0.90, 0.95])
    dinner_df = pd.DataFrame(dinner_series)
    all_rows = []
    for index in dinner_df.index:
        tmp_dict = dict(dinner_series[index[0]])
        tmp_dict['id'] = index[0]
        all_rows.append(tmp_dict)
    dinner_summary_df = pd.DataFrame(all_rows, columns = ['id', 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95])\
        .rename(columns = {0.05: '5%', 0.1: '10%', 0.25: '25%', 0.5: '50%', 0.75: '75%', 0.9: '90%', 0.95: '95%'})\
        .drop_duplicates().reset_index(drop = True)
    dinner_variability_df = dinner_summary_df.copy()

    for col in dinner_variability_df.columns:
        if col == 'id' or col == '50%':
            continue
        dinner_variability_df[col] = dinner_variability_df[col] - dinner_variability_df['50%']
    dinner_variability_df['50%'] = dinner_variability_df['50%'] - dinner_variability_df['50%']

    fig, ax = plt.subplots(1, 1, figsize = (10, 10), dpi=80)
    sns_plot = sns.distplot( dinner_variability_df['90%'] - dinner_variability_df['10%'] )
    ax.set(xlabel='Variation Distribution for Dinner (90% - 10%)', ylabel='Kernel Density Estimation')

    if export == True:
        dinner_variability_df.to_csv(out_path)

    return dinner_variability_df