"""Apply a function (or functions) across multiple columns

See source https://github.com/tidyverse/dplyr/blob/master/R/across.R
"""
from abc import ABC, abstractmethod
from typing import Any, Callable, Iterable, Mapping, Optional, Tuple, Union

import numpy
from pandas import DataFrame, Series
from pipda import register_func, evaluate_expr, evaluate_args, evaluate_kwargs
from pipda.utils import functype
from pipda.context import ContextBase
from pipda.symbolic import DirectRefAttr

from ..core.utils import (
    df_setitem, length_of, recycle_df, to_df, vars_select, get_option
)
from ..core.middlewares import CurColumn
from ..core.contexts import Context
from .tidyselect import everything

class Across:
    """Across object"""
    def __init__(
            self,
            data: DataFrame,
            cols: Optional[Iterable[str]] = None,
            fns: Optional[Union[
                Callable,
                Iterable[Callable],
                Mapping[str, Callable]
            ]] = None,
            names: Optional[str] = None,
            base0: Optional[bool] = None,
            args: Optional[Tuple[Any]] = None,
            kwargs: Optional[Mapping[str, Any]] = None
    ) -> None:
        cols = everything(data) if cols is None else cols
        if not isinstance(cols, (list, tuple)):
            cols = [cols]
        cols = data.columns[vars_select(data.columns, cols, base0=base0)]
        base0 = get_option('index.base.0', base0)

        fns_list = []
        if callable(fns):
            fns_list.append({'fn': fns})
        elif isinstance(fns, (list, tuple)):
            fns_list.extend(
                {'fn': fn, '_fn': i + int(not base0), '_fn1': i+1, '_fn0': i}
                for i, fn in enumerate(fns)
            )
        elif isinstance(fns, dict):
            fns_list.extend(
                {'fn': value, '_fn': key}
                for key, value in fns.items()
            )
        elif fns is not None:
            raise ValueError(
                'Argument `_fns` of across must be None, a function, '
                'a formula, or a dict of functions.'
            )

        self.data = data
        self.cols = cols
        self.fns = fns_list
        self.names = names
        self.args = args or ()
        self.kwargs = kwargs or {}

    def evaluate(
            self,
            context: Optional[Union[Context, ContextBase]] = None
    ) -> DataFrame:
        """Evaluate object with context"""
        if isinstance(context, Context):
            context = context.value

        if not self.fns:
            self.fns = [{'fn': lambda x: x}]
        ret = None
        for column in self.cols:
            for fn_info in self.fns:
                render_data = fn_info.copy()
                render_data['_col'] = column
                fn = render_data.pop('fn')
                name_format = self.names
                if not name_format:
                    name_format = (
                        '{_col}_{_fn}' if '_fn' in render_data
                        else '{_col}'
                    )

                name = name_format.format(**render_data)
                args = CurColumn.replace_args(self.args, column)
                kwargs = CurColumn.replace_kwargs(self.kwargs, column)
                if functype(fn) == 'plain':
                    value = fn(
                        self.data[column],
                        *evaluate_args(args, self.data, context),
                        **evaluate_kwargs(kwargs, self.data, context)
                    )
                else:
                    # use fn's own context
                    value = fn(
                        DirectRefAttr(self.data, column),
                        *args,
                        **kwargs,
                        _env='piping'
                    )._pipda_eval(self.data, context)

                if ret is None:
                    ret = to_df(value, name)
                elif length_of(ret) == 1:
                    ret, value = recycle_df(ret, value)
                    ret = df_setitem(ret, name, value)
                else:
                    ret = df_setitem(ret, name, value)
        return DataFrame() if ret is None else ret

class IfCross(Across, ABC):
    """Base class for IfAny and IfAll"""
    @staticmethod
    @abstractmethod
    def aggregate(values: Series) -> bool:
        """How to aggregation by rows"""

    def evaluate(self, context: Optional[ContextBase] = None) -> DataFrame:
        """Evaluate the object with context"""
        # Fill NA first and then do and/or
        # Since NA | True -> False for pandas
        return super().evaluate(context).apply(
            self.__class__.aggregate,
            axis=1
        ).astype(bool)

class IfAny(IfCross):
    """For calls from dplyr's if_any"""
    @staticmethod
    def aggregate(values: Series) -> bool:
        """How to aggregation by rows"""
        return values.fillna(False).astype(bool).any()

class IfAll(IfCross):
    """For calls from dplyr's if_all"""
    @staticmethod
    def aggregate(values: Series) -> bool:
        """How to aggregation by rows"""
        return values.fillna(False).astype(bool).all()

@register_func(
    context=Context.PENDING,
    verb_arg_only=True,
    summarise_prefers_input=True
)
def across(
        _data: DataFrame,
        *args: Any,
        _names: Optional[str] = None,
        _fn_context: Optional[Union[Context, ContextBase]] = Context.EVAL,
        _base0: Optional[bool] = None,
        **kwargs: Any
) -> DataFrame:
    """Apply the same transformation to multiple columns

    The original API:
    https://dplyr.tidyverse.org/reference/across.html

    Args:
        _data: The dataframe.
        *args: If given, the first 2 elements should be columns and functions
            apply to each of the selected columns. The rest of them will be
            the arguments for the functions.
        _names: A glue specification that describes how to name
            the output columns. This can use `{_col}` to stand for the
            selected column name, and `{_fn}` to stand for the name of
            the function being applied.
            The default (None) is equivalent to `{_col}` for the
            single function case and `{_col}_{_fn}` for the case where
            a list is used for _fns. In such a case, `{_fn}` is 0-based.
            To use 1-based index, use `{_fn1}`
        _base0: Indicating whether the columns are 0-based if selected
            by indexes. if not provided, will use
            `datar.base.get_option('index.base.0')`.
        _fn_context: Defines the context to evaluate the arguments for functions
            if they are plain functions.
            Note that registered functions will use its own context
        **kwargs: Keyword arguments for the functions

    Returns:
        A dataframe with one column for each column and each function.
    """
    if not args:
        args = (None, None)
    elif len(args) == 1:
        args = (args[0], None)
    _cols, _fns, *args = args
    _cols = evaluate_expr(_cols, _data, Context.SELECT)

    return Across(
        _data, _cols, _fns, _names, _base0, args, kwargs
    ).evaluate(_fn_context)

@register_func(context=Context.SELECT, verb_arg_only=True)
def c_across(
        _data: DataFrame,
        _cols: Optional[Iterable[str]] = None,
        _base0: Optional[bool] = None
) -> Series:
    """Apply the same transformation to multiple columns rowwisely

    Args:
        _data: The dataframe
        _cols: The columns
        _base0: Indicating whether the columns are 0-based if selected
            by indexes. if not provided, will use
            `datar.base.get_option('index.base.0')`.

    Returns:
        A series
    """
    if not _cols:
        _cols = everything(_data)

    _cols = vars_select(_data.columns.tolist(), _cols, base0=_base0)

    series = [_data.iloc[:, col] for col in _cols]
    return numpy.concatenate(series)

@register_func(
    context=None,
    extra_contexts={'args': Context.SELECT},
    verb_arg_only=True,
    summarise_prefers_input=True
)
def if_any(
        _data: DataFrame,
        *args: Any,
        _names: Optional[str] = None,
        _context: Optional[ContextBase] = None,
        _base0: Optional[bool] = None,
        **kwargs: Any
) -> Iterable[bool]:
    """Apply the same predicate function to a selection of columns and combine
    the results True if any element is True.

    See Also:
        [`across()`](datar.dplyr.across.across)
    """
    if not args:
        args = (None, None)
    elif len(args) == 1:
        args = (args[0], None)
    _cols, _fns, *args = args

    return IfAny(
        _data, _cols, _fns, _names, _base0, args, kwargs
    ).evaluate(_context)


@register_func(
    context=None,
    extra_contexts={'args': Context.SELECT},
    verb_arg_only=True,
    summarise_prefers_input=True
)
def if_all(
        _data: DataFrame,
        # _cols: Optional[Iterable[str]] = None,
        # _fns: Optional[Union[Mapping[str, Callable]]] = None,
        *args: Any,
        _names: Optional[str] = None,
        _context: Optional[ContextBase] = None,
        _base0: Optional[bool] = None,
        **kwargs: Any
) -> Iterable[bool]:
    """Apply the same predicate function to a selection of columns and combine
    the results True if all elements are True.

    See Also:
        [`across()`](datar.dplyr.across.across)
    """
    if not args:
        args = (None, None)
    elif len(args) == 1:
        args = (args[0], None)
    _cols, _fns, *args = args

    return IfAll(
        _data, _cols, _fns, _names, _base0, args, kwargs
    ).evaluate(_context)
