import bibliograph as bg
import pandas as pd


def _item_prefix_splitter(item_grp, prefixed_items):
    '''
    Dear Future Me: This function was refactored to ignore an
    argument "expected_prefixes" that was previously passed in.
    If you run into trouble you might have to revisit that.
    '''

    # get the labels for the item corresponding to the current
    item_syntax = prefixed_items.query('item_label == @item_grp.name')
    item_pfx_separator = bg.util.get_single_value(
        item_syntax,
        'item_prefix_separator'
    )

    # Get the default prefix and the set of possible prefixes
    # for this item from the entry syntax
    expected_prefixes = bg.util.get_single_value(
        item_syntax,
        'item_prefixes'
    )
    expected_prefixes = expected_prefixes.split()
    default_pfix = expected_prefixes[0]

    # add the item prefix separator onto the prefixes
    expected_prefixes = [
        p + item_pfx_separator for p in expected_prefixes
    ]
    default_pfix = default_pfix + item_pfx_separator

    # escape any regex metacharacters in the expected prefixes
    expected_prefixes = [
        bg.util.escape_regex_metachars(p) for p in expected_prefixes
    ]

    # make a regular expression that will match any of the
    # possible prefixes for this item
    prefixes_regex = '|'.join([f'^{p}' for p in expected_prefixes])
    has_no_pfix = ~item_grp.str.contains(prefixes_regex)

    # Add the default prefix onto unprefixed instances of this
    # item. Copy the group first so we don't mutate a grouped
    # object during iteration.
    pfixed = item_grp.copy()
    pfixed.loc[has_no_pfix] = pfixed.loc[has_no_pfix] \
                                    .apply(lambda x: default_pfix + x)

    # Split the group on the prefix separator and expand into a
    # dataframe, then return the dataframe
    return pfixed.str.split(item_pfx_separator, n=1, expand=True)


def _expand_grouped_entries(
    entry_grp,
    entry_syntax,
    item_separator,
    item_separator_regex,
    escaped_sep_regex,
    na_string_values,
    default_entry_prefix
):
    '''
    Takes a pandas.Series of stacked strings representing shorthand
    entries and expands them into their component items according to
    rules defined by the entry syntax.

    Helper function for parse_entries.

    Parameters
    ----------
    entry_grp : pandas.Series
        A group of entries generated by pandas.Series.groupby. This
        series has a multiindex generated by pandas.DataFrame.stack

    entry_syntax : pandas.DataFrame
        A dataframe containing item node types and item labels for each
        type of entry.

    item_separator : str
        A string separating items within an entry string.

    na_string_values : list-like
        Items within an entry whose values are in na_string_values will be
        replaced with pandas.NA

    Returns
    -------
    pandas.DataFrame
        A DataFrame the same length as the input with columns defined
        by the item_label column in entry_syntax.
    '''
    # Strip any trailing separators from each entry. This operation
    # creates a copy of the input group, so we're free to mutate the
    # copy
    expanded = entry_grp.str.rstrip(item_separator)

    if not pd.isna(entry_grp.name):
        grp_prefix = entry_grp.name
        slice_start = len(grp_prefix + item_separator)
        expanded = expanded.str.slice(start=slice_start)
    else:
        grp_prefix = default_entry_prefix

    # split on bare item separators and expand into a dataframe
    expanded = expanded.str.split(pat=item_separator_regex, expand=True)

    if (expanded[0] == '').any():
        # entries beginning with an item separator have self-descriptive
        # syntax
        raise NotImplementedError(
            'self-descriptive syntax is not yet implemented.'
        )

    else:
        # replace escaped item separators with the bare value
        expanded = expanded.replace(
            to_replace=escaped_sep_regex,
            value=item_separator,
            regex=True
        )

    # Replace NA values and empty strings with the first string in
    # na_string_values
    expanded = expanded.fillna(na_string_values[0])
    expanded = expanded.replace('', na_string_values[0])

    # column values were integers after str.split(pat, expand=True)
    # above but they need to be strings because string-valued item
    # labels are allowed
    expanded.columns = [str(c) for c in expanded.columns]

    # group prefixes are required to sort out entry links later, so
    # add an index level for the group prefix
    expanded['grp_prefix'] = [grp_prefix]*len(expanded)
    expanded = expanded.set_index('grp_prefix', append=True)

    # get item labels and node types for this entry prefix
    item_labels = entry_syntax.query('entry_prefix == @grp_prefix')

    # items with no node type in the entry syntax are prefixed to
    # indicate which node type they correspond to
    item_is_prefixed = item_labels['item_node_type'].isna()

    if item_is_prefixed.any():

        prefixed_items = item_labels.loc[item_is_prefixed]
        labels_of_prefixed_items = prefixed_items['item_label'].array

        # stack the prefixed items into a series
        disagged = expanded[labels_of_prefixed_items].stack()

        # Split the prefixes off of the stacked items and expand into a
        # dataframe
        disagged = disagged.groupby(level=3).apply(
            _item_prefix_splitter,
            prefixed_items
        )

        # drop the item labels from the multiindex so the disaggregated
        # items align with the index of the entry group
        disagged.index = disagged.index.droplevel(3)

        # pivot the disaggregated items to create a dataframe with
        # columns for each item prefix
        disagged = disagged.pivot(columns=0)
        disagged.columns = disagged.columns.get_level_values(1)

        # get labels of items that are not prefixed and present in this
        # dataset
        unprefixed_item_labels = [
            label for label in item_labels['item_label']
            if label.isdigit()
            and label in expanded.columns
            and label not in labels_of_prefixed_items
        ]

        # select only the unprefixed item labels
        expanded = expanded[unprefixed_item_labels]
        # concatenate the unprefixed and prefixed items and return
        return pd.concat([expanded, disagged], axis='columns')

    else:
        # there were no prefixed items, so return the expanded items
        return expanded


def parse_entries(
    entries,
    entry_syntax,
    item_separator,
    default_entry_prefix,
    space_char,
    na_string_values,
    s_d_delimiter
):
    '''
    Take a stacked pandas.Series of shorthand entry strings and expands
    them into a dataframe with the following columns:
        ['csv_row',
         'csv_col',
         'entry_prefix',
         'item_label',
         'string',
         'node_type',
         'item_link_type',
         'node_tags']

    Parameters
    ----------
    entries : pandas.Series
        A series of strings interpreted as shorthand entries. Expects a
        multiindex with the row and column index of each entry in a csv
        file.

    entry_syntax : pandas.DataFrame
        A dataframe containing item node types and item labels for each
        type of entry.

    item_separator : str
        A string separating items within an entry string.

    default_entry_prefix : str
        Entries with no prefix will be interpreted as having the default
        prefix.

    space_char : str
        A string value that should be replaced with a single space
        unless escaped with a backslash.

    na_string_values : str or iterable of strings
        A single string or set of string values which represent missing
        items in an entry.

    s_d_delimiter : str
        A string-valued delimiter separating elements in a
        self-descriptive shorthand entry

    Returns
    -------
    pandas.DataFrame
        Entries expanded according to the entry syntax
    '''

    if entries.empty:
        raise ValueError('entries cannot be empty')

    # escape any regex metacharacters in the item separator so we can
    # use it in regular expressions
    regex_item_separator = bg.util.escape_regex_metachars(item_separator)

    # Use a regular expression to split tags off the input strings. Tags
    # are separated from entries by an item separator followed by a
    # space.
    tag_sep_regex = r"(?<!\\)[{}][ ]".format(']['.join(regex_item_separator))
    entries = entries.str.split(pat=tag_sep_regex, expand=True)

    if len(entries.columns) == 1:
        entries = entries.rename(columns={0: 'string'})
        entries['node_tags'] = pd.NA
    else:
        entries = entries.rename(columns={0: 'string', 1: 'node_tags'})

    # Use a regular expression to get entry prefixes
    # (null if no prefix)
    prefixes = [
        bg.util.escape_regex_metachars(p) for p in
        entry_syntax['entry_prefix'].dropna().drop_duplicates()
    ]
    prefixes = prefixes + [regex_item_separator]
    prefixes_regex = '|'.join(prefixes)
    prefixes_regex = '^({})(?={})'.format(prefixes_regex, regex_item_separator)

    entries['entry_prefix'] = entries['string'].str.extract(prefixes_regex)

    # regular expressions to match bare and escaped item separators
    item_separator_regex = r"(?<!\\)[{}]".format(
        ']['.join(regex_item_separator)
    )
    escaped_sep_regex = fr"(\\{regex_item_separator})"

    # no_prefix = entries['entry_prefix'].isna()
    # entries.loc[no_prefix, 'entry_prefix'] = default_entry_prefix
    # Group entries by prefix and expand them
    expanded = entries['string'].groupby(
        by=entries['entry_prefix'],
        dropna=False,
        group_keys=False
    )
    expanded = expanded.apply(
        _expand_grouped_entries,
        entry_syntax,
        item_separator,
        item_separator_regex,
        escaped_sep_regex,
        na_string_values,
        default_entry_prefix
    )

    # expanded has a multiindex whose levels correspond to csv rows,
    # csv columns, and entry prefixes

    # Replace any empty strings with null values
    expanded = expanded.mask(expanded == '', pd.NA)

    # Regular expressions to match bare and escaped space placeholders
    regex_space_char = bg.util.escape_regex_metachars(space_char)
    space_plchldr_regex = r"(?<!\\)({})".format(regex_space_char)
    escaped_space_plchldr_regex = fr"(\\{regex_space_char})"

    # Replace space placeholders with spaces in the expanded items
    expanded = expanded.replace(
        to_replace=space_plchldr_regex,
        value=' ',
        regex=True
    )
    # Replace escaped space placeholders with bare placeholders
    expanded = expanded.replace(
        to_replace=escaped_space_plchldr_regex,
        value=regex_space_char,
        regex=True
    )
    # Stack the expanded items. Stacking creates a series whose values
    # are the string values of every item in the input entries and
    # whose index levels are
    #       csv rows, csv columns, entry prefix, item label
    expanded = expanded.stack()

    # Locate any items that came out of an entry with self-descriptive
    # syntax.
    item_is_s_d = (expanded.index.get_level_values(2) == regex_item_separator)
    s_d_items_exist = item_is_s_d.any()

    if s_d_items_exist:

        # S-d entry node types are first items in s-d entries, so locate
        # those
        item_is_first_in_entry = (
            expanded.index.get_level_values(3) == '0'
        )
        s_d_entry_node_types = expanded.loc[
            item_is_s_d & item_is_first_in_entry,
            :
        ]
        entry_is_s_d = (entries['entry_prefix'] == regex_item_separator)
        entries.loc[entry_is_s_d, 'entry_prefix'] = s_d_entry_node_types.array

        # You can't use syntactic prefixes as node types in s-d entries
        node_type_collisions = s_d_entry_node_types.isin(
            entry_syntax['entry_prefix']
        )
        if node_type_collisions.any():
            collisions = s_d_entry_node_types.loc[node_type_collisions]
            raise ValueError(
                'Self-descriptive entries cannot use entry node types '
                'that are also entry prefixes in the entry syntax.\n'
                '            Found the following syntactic entry '
                'prefixes in self-descriptive entries:\n'
                '            {}'.format(list(collisions))
            )

        # Self-descriptive entries might have different numbers of
        # items, leading the parser to introduce missing values not
        # present in the input data and assign them the first
        # na_string_value. Locate items whose string values are not the
        # first na_string_value and drop them.
        item_is_not_na_value_0 = (expanded != na_string_values[0])
        s_d_items_to_keep = (
            item_is_s_d & item_is_not_na_value_0 & ~item_is_first_in_entry
        )

        # Drop the s-d entry node types and missing values created by
        # the parser and separate the syntactic entries from the
        # s-d entries
        expanded_s_d_items = expanded.loc[s_d_items_to_keep]
        expanded_syntactic_items = expanded.loc[~item_is_s_d]

        # Get the number of items in each s-d entry aside from the
        # entry node type
        s_d_entry_lengths = expanded_s_d_items.groupby(level=[0, 1])
        s_d_entry_lengths = s_d_entry_lengths.apply(
            lambda x: len(x)
        )

        # Self-descriptive items must have exactly two delimiters,
        # one between the link type and the node type and the other
        # between the node type and the string value
        num_delimiters = expanded_s_d_items.map(
            lambda x: x.count(s_d_delimiter)
        )
        if (num_delimiters != 2).any():
            raise ValueError(
                'Self-descriptive entries must have exactly two '
                'delimiters in every item.'
            )

        # Separate the node types, link types, and string values in the
        # s-d items
        expanded_s_d_items = expanded_s_d_items.str.split(
            s_d_delimiter,
            expand=True
        )

        # Create an index for the s-d items that includes the entry
        # node types (which might be prefixes from the entry syntax)
        s_d_entry_node_types = [
            v
            for i in range(len(s_d_entry_lengths))
            for v in ([s_d_entry_node_types.iloc[i]]*s_d_entry_lengths.iloc[i])
        ]

        # Make an index for the s-d items that has the entry node types
        # where the entry prefixes should appear
        s_d_idx = expanded_s_d_items.index
        s_d_item_labels = s_d_idx.get_level_values(3)
        s_d_idx = pd.MultiIndex.from_arrays((
            s_d_idx.get_level_values(0),
            s_d_idx.get_level_values(1),
            s_d_entry_node_types,
            s_d_item_labels
        ))
        s_d_idx.names = [None, None, 'grp_prefix', None]

        # Create a map from s-d entry node types (which could be entry
        # prefixes from the entry syntax) and item labels to item node
        # types and item link types
        s_d_item_types = pd.DataFrame({
            'entry_prefix': s_d_entry_node_types,
            'item_label': s_d_item_labels,
            'node_type': expanded_s_d_items[1].array,
            'link_type': expanded_s_d_items[0].array
        })
        s_d_item_types = s_d_item_types.drop_duplicates()
        s_d_item_types = s_d_item_types.set_index('entry_prefix', drop=True)
        s_d_item_types = s_d_item_types.set_index(
            'item_label',
            append=True,
            drop=True
        )

        # Include the s-d data in the expanded items series
        expanded = pd.concat([
            pd.Series(expanded_s_d_items[2].array, index=s_d_idx),
            expanded_syntactic_items
        ])

    # Convert expanded to a dataframe so we can concatenate it with
    # columns
    # expanded = pd.DataFrame(expanded, columns=['string'])

    # Create a map from entry prefixes and item labels to item node
    # types and item link types
    item_label_idx = pd.MultiIndex.from_arrays(
        (entry_syntax['entry_prefix'], entry_syntax['item_label'])
    )
    item_types = pd.DataFrame(
        {'node_type': entry_syntax['item_node_type'].array,
         'link_type': entry_syntax['item_link_type'].array},
        index=item_label_idx
    )

    # Add self-descriptive item types if they exist
    if s_d_items_exist:
        item_types = pd.concat([item_types, s_d_item_types])

    # Get an index of entry prefixes and item labels from the expanded
    # items.
    entry_item_idx = expanded.index.droplevel([0, 1])
    # Map the item node and link types, set the index to conform to the
    # index for the expanded items.
    item_types = item_types.loc[entry_item_idx].set_index(expanded.index)

    # Concatenate the item node and link types onto the item strings.
    expanded = pd.concat([expanded.rename('string'), item_types], axis=1)
    # expanded = _set_StringDtype(expanded)

    # join tag strings back onto values in the 'string' column to recover
    # the original input strings
    where_tags = entries['node_tags'].notna()
    entries.loc[where_tags, 'string'] = entries.loc[where_tags, 'string'] \
        + item_separator + ' ' + entries.loc[where_tags, 'node_tags']

    # Make a map between entry prefixes and node types
    node_type_map = entry_syntax.loc[:, ['entry_prefix', 'entry_node_type']]
    node_type_map = node_type_map.drop_duplicates()
    node_type_map = pd.Series(
        node_type_map['entry_node_type'].array,
        index=node_type_map['entry_prefix'].array
    )
    # Add a row in the node type map for entries with no prefix
    default_entry_type = bg.util.get_single_value(
        entry_syntax.query('entry_prefix == @default_entry_prefix'),
        'entry_node_type'
    )
    node_type_map[pd.NA] = default_entry_type
    # Add self-descriptive node types to the node type map
    if s_d_items_exist:

        s_d_node_types = pd.Series(s_d_item_types.index.get_level_values(0))
        s_d_node_types = s_d_node_types.drop_duplicates()

        # If there are syntactic entry prefixes in the self-descriptive
        # node types, don't add new map data for them
        s_d_node_types = s_d_node_types.loc[
            ~s_d_node_types.isin(node_type_map.index)
        ]

        # Self-descriptive entry node "prefixes" are their own node
        # types
        s_d_node_types = pd.Series(s_d_node_types.array, index=s_d_node_types)

        node_type_map = pd.concat([node_type_map, s_d_node_types])

    # Normalize null types
    entries = entries.fillna(pd.NA)

    # Create a column in entries for node types
    entries['node_type'] = node_type_map[entries['entry_prefix']].array
    # entries['node_type'] = entries['node_type'].astype(pd.StringDtype())
    # can't use pd.StringDtype() throughout because it currently doesn't allow
    # construction with null types other than pd.NA. This will likely change
    # soon (https://github.com/pandas-dev/pandas/pull/41412)

    # Add levels to the multiindex of the input entries to conform with
    # the index of the expanded items.
    entries = entries.set_index('entry_prefix', append=True, drop=True)
    entries['item_label'] = pd.NA
    entries = entries.set_index('item_label', append=True, drop=True)

    # Concatenate expanded items with the entries
    entries = pd.concat([expanded, entries]).sort_index().fillna(pd.NA)

    return entries


def _expand_csv_items(group, delimiters):
    '''
    THIS FUNCTION MUTATES ITS FIRST ARGUMENT

    Takes item strings grouped by item label and splits the strings if
    given a list delimiter for the group.
    '''

    item_label_id = group.name

    if pd.isna(item_label_id):
        return group

    delimiter = delimiters.loc[item_label_id]

    if pd.isna(delimiter):
        return group

    else:
        group.loc[:, 'string'] = group['string'].str.split(delimiter)
        return group
