from collections import defaultdict
from typing import Union, List, Dict, Callable, Mapping, Iterable

import dask.dataframe as dd
import numpy as np
import pandas as pd
from pandas.core.groupby import SeriesGroupBy


def get_agg_func(keyword: Union[str, Callable], use_dask=False) -> Union[str, Callable, dd.Aggregation]:
    """

    Args:
        keyword (str):
        use_dask (bool): Whether to create a dd.Aggregation

    Returns:
        func (callable): a callable function, pandas aggregator func name, or a Dask Aggregation.
    """
    if callable(keyword) or isinstance(keyword, dd.Aggregation):
        return keyword

    elif keyword == "unique" and use_dask:
        # get unique values (in a list-like np.array) from each groupby key
        func = concat_unique_dask_agg()
    elif keyword == "unique" and not use_dask:
        func = concat_uniques

    elif keyword == "concat":
        # Concatenate values into list
        func = concat

    else:
        # Any other aggregation keywords or callable function
        func = keyword

    return func


def get_multi_aggregators(agg: str, agg_for: Dict[str, Union[str, Callable, dd.Aggregation]] = None, use_dask=False) \
    -> Mapping[str, Union[str, dd.Aggregation]]:
    """

    Args:
        agg ():
        agg_for ():
        use_dask ():

    Returns:

    """
    if agg_for is None:
        agg_for = {}

    col2func = {col: get_agg_func(keyword, use_dask=use_dask) for col, keyword in agg_for.items()}
    col_aggregators = defaultdict(lambda: get_agg_func(agg, use_dask=use_dask), col2func)

    return col_aggregators


def concat_unique_dask_agg() -> dd.Aggregation:
    def chunk(s: pd.Series) -> pd.Series:
        '''
        The function applied to the individual partition (map)
        '''

        def to_list(x: Union[str, List, np.ndarray]) -> List:
            if isinstance(x, str):
                return [x]
            elif isinstance(x, np.ndarray):
                return x
            elif isinstance(x, Iterable):
                if any(isinstance(a, Iterable) for a in x):
                    return list(set(np.hstack(x)))
                else:
                    return list(set(x))
            else:
                return [x]

        return s.apply(to_list)

    def agg(s: SeriesGroupBy) -> pd.Series:
        '''
        The function which will aggregate the result from all the partitions(reduce)
        '''
        s = s._selected_obj
        return s.groupby(level=list(range(s.index.nlevels)), group_keys=True).apply(
            lambda li: np.hstack(li) if isinstance(li, Iterable) and len(li) else None)

    def finalize(s) -> pd.Series:
        '''
        The optional functional that will be applied to the result of the agg_tu functions
        '''
        return s.apply(lambda arr: np.unique(arr[~pd.isna(arr)]))

    func = dd.Aggregation(
        name='unique',
        chunk=chunk,
        agg=agg,
        finalize=finalize
    )
    return func


def merge_concat(a: Union[str, None, Iterable], b: Union[str, None, Iterable]) -> Union[np.ndarray, str, None]:
    """
    Used as function in pd.combine() or dd.combine()
    Args:
        a (Union[str,None,Iterable]): cell value in a pd.Series
        b (Union[str,None,Iterable]): cell value in a pd.Series

    Returns:
        combined_value (Union[np.ndarray, str, None])
    """
    a_isna = pd.isna(a)
    b_isna = pd.isna(b)
    if a_isna is True or (isinstance(a_isna, Iterable) and all(a_isna)):
        return b
    elif b_isna is True or (isinstance(b_isna, Iterable) and all(b_isna)):
        return a
    elif isinstance(a, str) and isinstance(b, str):
        if a == b:
            return a
        return np.array([a, b])
    elif not isinstance(a, Iterable) and isinstance(b, Iterable):
        return np.hstack([[a], b])
    elif isinstance(a, Iterable) and not isinstance(b, Iterable):
        return np.hstack([a, [b]])
    elif isinstance(a, Iterable) and isinstance(b, Iterable):
        return np.hstack([a, b])
    else:
        return b


def concat_uniques(series: pd.Series) -> Union[str, List, np.ndarray, None]:
    """ An aggregation custom function to be applied to each column of a groupby
    Args:
        series (pd.Series): Entries can be either a string or a list of strings.
    Returns:
        unique_values
    """
    series = series.dropna()
    if series.empty:
        return None

    is_str_idx = series.map(type) == str

    if series.map(lambda x: isinstance(x, Iterable)).any():
        if (is_str_idx).any():
            # Convert mixed dtypes to lists
            series.loc[is_str_idx] = series.loc[is_str_idx].map(lambda s: [s] if len(s) else None)
        return np.unique(np.hstack(series))

    elif is_str_idx.any():
        concat_str = series.astype(str).unique()
        if len(concat_str):  # Avoid empty string
            return concat_str

    else:
        return series.tolist()

def concat(series: pd.Series) -> Union[str, List, np.ndarray, None]:
    """
    Args:
        series (pd.Series): Entries can be either a string or a list of strings.
    """
    series = series.dropna()
    if series.empty:
        return

    is_str_idx = series.map(type) == str
    if series.map(lambda x: isinstance(x, Iterable)).any():
        if (is_str_idx).any():
            # Convert mixed dtypes to lists
            series.loc[is_str_idx] = series.loc[is_str_idx].map(lambda s: [s] if len(s) else None)
        return np.hstack(series)

    elif is_str_idx.any():
        concat_str = series.astype(str).tolist()
        if len(concat_str):  # Avoid empty string
            return concat_str

    else:
        return series.tolist()