from collections.abc import Iterable

from typing import Union, Dict, List, Tuple, Optional, Set

import numpy as np

import pandas as pd

from dask import dataframe as dd

from logzero import logger

def has_iterables(series: Union[dd.Series, pd.Series], n=10) -> bool:

    """

    Check whether any element in series is an Iterable, e.g. list, set, or np.ndarray.

    Args:

        series ():

        n (int): number of elements to test

    Returns:

        bool

    """

    if isinstance(series, (dd.Series, dd.Index)):

        try:

            values = series.head(n=n, npartitions=-1)

        except:

            values = series.head(n=n)

    elif isinstance(series, pd.Series):

        values = series.head(n=n)

    elif isinstance(series, pd.Index):

        values = series[:n]

    else:

        return False

    is_iterables = values.map(lambda x: not isinstance(x, str) and isinstance(x, Iterable)).any()

    return is_iterables

def match_iterable_keys(left: Union[dd.Series, pd.Series], right: Union[dd.Series, pd.Series]) \

    -> Tuple[dd.Series, dd.Series]:

    left_iterables = has_iterables(left)

    right_iterables = has_iterables(right)

    def _list_to_key(list_values: List[str], possibilities: Set[str]) -> Optional[str]:

        if list_values is None:

            return None

        elif len(list_values) == 1:

            return list_values[0]

        for key in list_values:

            if key in possibilities:

                return key

    def _list_to_list(list_values: List[str], multi_possibilities: List[Set[str]]) -> Optional[str]:

        if list_values is None:

            return None

        elif len(list_values) == 1:

            return list_values[0]

        for possibilities in multi_possibilities:

            if isinstance(possibilities, set) and list_values is not None:

                match = possibilities.intersection(list_values)

                if len(match):

                    return "|".join(sorted(match))

    left_on, right_on = left, right

    if left_iterables and not right_iterables:

        possibilities = set(right.dropna()) if isinstance(right, (pd.Series, pd.Index)) else set(

            right.dropna().compute())

        left_on = left.map(lambda values: _list_to_key(values, possibilities))

    elif not left_iterables and right_iterables:

        possibilities = set(left.dropna()) if isinstance(left, (pd.Series, pd.Index)) else set(

            left.dropna().compute())

        right_on = right.map(lambda values: _list_to_key(values, possibilities))

    elif left_iterables and right_iterables:

        right_possibilities = right.map(lambda x: set(x) if isinstance(x, Iterable) and not isinstance(x, str) else x)

        if isinstance(right_possibilities, (dd.Series, dd.Index)):

            right_possibilities = right_possibilities.compute()

        left_on = left.map(lambda values: _list_to_list(values, right_possibilities))

        left_possibilities = left.map(lambda x: set(x) if isinstance(x, Iterable) and not isinstance(x, str) else x)

        if isinstance(left_possibilities, (dd.Series, dd.Index)):

            left_possibilities = left_possibilities.compute()

        right_on = right.map(lambda values: _list_to_list(values, left_possibilities))

    return left_on, right_on

def drop_duplicate_columns(df: Union[pd.DataFrame, dd.DataFrame]) -> Union[pd.DataFrame, dd.DataFrame]:

    """

    Args:

        df:

    """

    if df.columns.duplicated().any():

        _, i = np.unique(df.columns, return_index=True)

        df = df.iloc[:, i]

    return df

def filter_rows(df: pd.DataFrame, filters: Union[str, Dict[str, List]], uncased=False):

    """

    Args:

        df (pd.DataFrame):

        filters (str or dict):

            Either a pandas query expression or a dict of column names for keys and matching values.

        uncased (bool): Default False.

            Whether to match case in pd.Series.str.contains if filters is a dict of values.

    Returns:

    """

    num_samples = df.shape[0]

    if filters is None:

        return df

    elif isinstance(filters, str):

        df = df.query(filters)

    elif isinstance(filters, dict):

        for col, values in filters.items():

            if col not in df.columns:

                logger.warn("Filter key `", col, "` must be in one of ", df.columns)

                continue

            if isinstance(values, list):

                values = {val.upper() for val in values}

                if uncased:

                    df = df.loc[df[col].str.upper().isin(values)]

                else:

                    df = df.loc[df[col].isin(values)]

            elif isinstance(values, str):

                df = df.loc[df[col].str.contains(values, case=not uncased)]

            else:

                df = df.loc[df[col] == values]

    if isinstance(num_samples, int):

        logger.info(f'Removed {num_samples - df.shape[0]} rows from query: {filters}')

    return df