# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/04-updates-and-cleaning.ipynb (unless otherwise specified).

__all__ = ['filter_cols_for_one_to_many', 'apply_updates', 'check_output_col_values', 'apply_output_col_types',
           'check_and_apply_output_defs']

# Cell
import json
import pandas as pd

import os
import typing
from warnings import warn

from powerdict import construct

# Cell
def filter_cols_for_one_to_many(df, cols=['longitude', 'latitude']):
    df_one_to_many_filt = df[df[cols].astype(str).agg(''.join, axis=1).str.contains(', ')]
    return df_one_to_many_filt

# Cell
def apply_updates(df, updates_data_dir):
    update_files = [f for f in os.listdir(updates_data_dir) if f.replace('.json', '') in df.columns]

    for update_file in update_files:
        with open(f'{updates_data_dir}/{update_file}', 'r') as f:
            update_dict = json.load(f)

        update_col = update_file.replace('.json', '')
        update_values = {k: v['new_value'] for k, v in update_dict.items()}

        df = construct.update_df_col(df, update_col, update_values)

    return df

# Cell
def check_output_col_values(df, output_col, output_defs):
    output_col_def = output_defs[output_col]

    if isinstance(output_col_def, list):
        output_col_elems = list()

        for elem in df[output_col].to_list():
            if ', ' in elem:
                output_col_elems += elem.split(', ')
            else:
                output_col_elems += [elem]

        unexpected_col_elems = sorted(list(set(output_col_elems) - set(output_col_def)))

        if len(unexpected_col_elems) > 0:
            warn(f"The following list elements were not expected in {output_col}: {', '.join(unexpected_col_elems)}")

    return

# Cell
def apply_output_col_types(df, output_col, output_defs):
    output_col_def = output_defs[output_col]

    if isinstance(output_col_def, list):
        dtype = type(output_col_def[0])
    else:
        dtype = output_col_def

    df[output_col] = df[output_col].astype(dtype)

    return df

# Cell
def check_and_apply_output_defs(df, definitions_dir):
    with open(f'{definitions_dir}/outputs.json', 'r') as f:
        output_defs = json.load(f)

    for output_col in output_defs.keys():
        if output_col not in df.columns:
            warn(f'{output_col} was not found in the output dataset')
        else:
            check_output_col_values(df, output_col, output_defs)
            df = apply_output_col_types(df, output_col, output_defs)

    return df