"""
Utility functions for Pandas.
"""

import re
from functools import lru_cache
from typing import Iterable

import numpy as np
import pandas as pd

from ..typings import ColumnName


def add_col_if_not_exists(
    df: pd.DataFrame,
    col_name: ColumnName | Iterable[ColumnName],
    default_value=np.nan,
    dtype=None,
) -> None:
    """Add a column to a DataFrame if it doesn't exist yet.

    Args:
        df (pd.DataFrame): DataFrame to add the column to.
        col_name (ColumnName): Name of the column to add.
        default_value: Default value to fill the column with.
        dtype: Data type of the column. If None, the data type is not changed and
            will be purely handled by NumPy.
    """
    if isinstance(col_name, str):
        col_name = [col_name]
    existing_df_cols = set(df.columns)
    cols_to_add = set(col_name) - existing_df_cols
    for col in cols_to_add:
        df[col] = default_value
        if dtype is None:
            continue
        # Indices where nans
        nan_indices = df[col].isna()
        # Convert column to the desired type and overwrite any values which
        # were NaN before the conversion, with nan
        df.loc[nan_indices, col] = np.nan
        df[col].infer_objects()


def extract_rows(df: pd.DataFrame, condition) -> pd.DataFrame:
    """Extract rows matching a condition from a DataFrame and return
    them as a new DataFrame (subset). The original DataFrame is modified
    in place, matching rows are removed from it, and the index is reset.
    The subset has the same schema as the original DataFrame.

    Args:
        df (pd.DataFrame): DataFrame to extract rows from.
        condition: Condition to match rows against, e.g. df["column"].ne(1).

    Returns:
        pd.DataFrame: Subset of the original DataFrame, containing only the
            rows that matched the condition.
    """
    indices = df.index[condition].tolist()
    subset = df.loc[indices]
    df.drop(indices, inplace=True)
    # DO NOT RESET INDEX! This way excluded rows and non-excluded rows will not
    # have overlapping indices, which is useful for keeping references to the
    # original DataFrame.
    return subset


def group_into_set(df: pd.DataFrame, agg_col: ColumnName) -> pd.DataFrame:
    """Group a DataFrame by all columns except one and aggregate this
    specific column into a set.
    """
    other_cols = [col for col in df.columns if col != agg_col]
    return df.groupby(other_cols)[agg_col].agg(set).reset_index()


def to_set(value: str | float, separator_regex: re.Pattern) -> set[str] | float:
    """Wrapper around _to_set that returns a copy of the set, so that
    the resulting object is guaranteed to be unique in memory, while increasing
    performance thanks to caching.

    Can be used with apply() to convert values in a Pandas DataFrame column into
    sets.
    """
    result = _to_set(value, separator_regex)
    return result.copy() if isinstance(result, set) else result


@lru_cache(maxsize=256)
def _to_set(value: str | float, separator_regex: re.Pattern) -> set[str] | float | None:
    """Split a string by a separator and return a set of unique values.

    This is a high performance function that uses caching for performance.
    It can also handle NaN values, which are ignored.

    ATTENTION! Because caching is used, this function will return the same
    object for the same input! To convert the resulting set into an unique
    object, create a copy of the outputted set.

    Args:
        value: String to split.
        separator_regex: Regular expression used to split the string.

    Examples:
        >>> len(to_set("ES, FR, IT"))
        3
    """
    try:
        splitted = separator_regex.split(value)
    except TypeError as ex:
        if isinstance(value, float) or value is None:
            # Ignore NaN (which is a float). Return the NaN value instead
            # of an empty set for performance and memory reasons.
            return value
        raise ex
    return set(map(str.strip, splitted))