# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/03-database-construction.ipynb (unless otherwise specified).

__all__ = ['load_source_definitions', 'identify_primary_keys', 'check_source_for_disallowed_cols',
           'check_sources_for_disallowed_cols', 'identify_inputs_for_output_cols', 'get_input_key_to_attr',
           'get_primary_key_to_attr', 'update_df_col', 'add_attr_cols', 'add_key_cols', 'construct_output_df']

# Cell
import json
import pandas as pd

import os
from warnings import warn

# Cell
def load_source_definitions(definitions_dir):
    source_definitions = dict()
    source_def_filenames = [f for f in os.listdir(definitions_dir) if '.json' in f and f != 'outputs.json']

    for filename in source_def_filenames:
        source_name = filename.replace('.json', '')

        with open(f'{definitions_dir}/{filename}', 'r') as f:
            source_definitions[source_name] = json.load(f)

    return source_definitions

# Cell
def identify_primary_keys(source_definitions):
    primary_keys = []

    for source in source_definitions.values():
        primary_keys += list(source['key_map'].keys())

    primary_keys = sorted(set(primary_keys))

    return primary_keys

# Cell
def check_source_for_disallowed_cols(attr_cols, outputs):
    output_cols_to_be_added = [attr_col['output_col'] for attr_col in attr_cols.values()]
    output_cols_allowed = outputs.keys()

    disallowed_output_cols = list(set(output_cols_to_be_added) - set(output_cols_allowed))

    assert len(disallowed_output_cols)==0, f"The following columns are not allowed in the output dataset: {', '.join(disallowed_output_cols)}"

def check_sources_for_disallowed_cols(source_definitions, definitions_dir):
    with open(f'{definitions_dir}/outputs.json', 'r') as f:
        outputs = json.load(f)

    for source, definition in source_definitions.items():
        check_source_for_disallowed_cols(definition['attr_cols'], outputs)

    return

# Cell
def identify_inputs_for_output_cols(source_definitions):
    keys_to_keep = ['output_col', 'output_rank']
    output_cols = list()
    inputs_for_output_cols = dict()

    for source, definition in source_definitions.items():
        for input_col, value in definition['attr_cols'].items():
            output_cols += [(source, input_col, value['output_col'], value['output_rank'])]

    df_output_cols = pd.DataFrame(output_cols, columns=['source', 'input_col', 'output_col', 'output_rank'])

    for output_col in df_output_cols['output_col']:
        inputs_for_output_cols[output_col] = (df_output_cols
                                              .query('`output_col`==@output_col')
                                              .sort_values('output_rank', ascending=False)
                                              [['source', 'input_col']]
                                              .apply(tuple, axis=1)
                                              .to_list()
                                             )
    return inputs_for_output_cols

# Cell
def get_input_key_to_attr(df_input, key_input_col, input_col):
    df_key_input_non_duped = df_input[[key_input_col, input_col]].drop_duplicates()

    one_to_many_input_keys = list(df_key_input_non_duped
                                  [df_key_input_non_duped[key_input_col].duplicated(keep=False)]
                                  .sort_values(key_input_col)
                                  [key_input_col]
                                  .unique()
                                 )

    warn(f"The following {key_input_col} input keys were not unique w.r.t. {input_col}, a one to many relationship will therefore be used: {', '.join(one_to_many_input_keys)}\n")

    # Creating one-to-one mapping
    input_key_to_attr = (df_key_input_non_duped
                         [~df_key_input_non_duped['bmu_root'].isin(one_to_many_input_keys)]
                         .set_index(key_input_col)
                         [input_col]
                         .to_dict()
                        )

    # Creating one-to-many mapping
    for input_key in one_to_many_input_keys:
        many_inputs = (df_key_input_non_duped
                       .query(f'{key_input_col}=="{input_key}"')
                       [input_col]
                       .to_list()
                      )

        input_key_to_attr[input_key] = ', '.join([str(elem) for elem in many_inputs])

    return input_key_to_attr

# Cell
def get_primary_key_to_attr(source, input_col, source_definitions, raw_data_dir):
    filename = source_definitions[source]['filename']
    key_input_col = source_definitions[source]['key_input_col']
    key_map = source_definitions[source]['key_map']
    value_map = source_definitions[source]['attr_cols'][input_col]['value_map']

    df_input = pd.read_csv(f"{raw_data_dir}/{filename}")

    if value_map is not None:
        df_input[input_col].isin(value_map.keys()).mean()==1, f'A mapping was provided for {input_col} but did not contain all of the values found within it'
        df_input[input_col] = df_input[input_col].map(value_map)

    input_key_to_attr = get_input_key_to_attr(df_input, key_input_col, input_col)

    primary_key_to_attr = (pd.Series(key_map)
                           .map(input_key_to_attr)
                           .to_dict())

    return primary_key_to_attr

# Cell
def update_df_col(df, update_col, update_dict):
    s_update = pd.Series(df.index.map(update_dict), df.index).dropna()
    df.loc[s_update.index, update_col] = s_update.values

    return df

def add_attr_cols(df, source_definitions, raw_data_dir):
    inputs_for_output_cols = identify_inputs_for_output_cols(source_definitions)

    for output_col, inputs in inputs_for_output_cols.items():
        for (source, input_col) in inputs:
            primary_key_to_attr = get_primary_key_to_attr(source, input_col, source_definitions, raw_data_dir)
            df = update_df_col(df, output_col, primary_key_to_attr)

    return df

# Cell
def add_key_cols(df, source_definitions):
    for source in source_definitions.keys():
        key_map = source_definitions[source]['key_map']
        key_output_col = source_definitions[source]['key_output_col']
        df = update_df_col(df, key_output_col, key_map)

    return df

# Cell
def construct_output_df(definitions_dir, raw_data_dir):
    source_definitions = load_source_definitions(definitions_dir)
    primary_keys = identify_primary_keys(source_definitions)

    df = (pd.DataFrame(index=primary_keys)
          .pipe(add_key_cols, source_definitions)
          .pipe(add_attr_cols, source_definitions, raw_data_dir)
         )

    df.index.name = 'osuked_id'

    return df