__all__ = ["GroupBy"]

import warnings
from enum import IntEnum

import numpy as np
import riptide_cpp as rc

from .rt_enum import GB_FUNC_COUNT, GB_FUNCTIONS, GB_STRING_ALLOWED, CategoryMode
from .rt_groupbykeys import GroupByKeys
from .rt_groupbyops import GroupByOps
from .rt_numpy import (
    _groupbycalculateall,
    _groupbycalculateallpack,
    empty_like,
    groupbypack,
    sortinplaceindirect,
)
from .rt_timers import GetTSC
from .rt_utils import bool_to_fancy, mbget

#'''
# Series.abs()	Return an object with absolute value taken�only applicable to objects that are all numeric.
# Series.all([axis, bool_only, skipna, level])	Return whether all elements are True over requested axis
# Series.any([axis, bool_only, skipna, level])	Return whether any element is True over requested axis
# Series.autocorr([lag])	Lag-N autocorrelation
# Series.between(left, right[, inclusive])	Return boolean Series equivalent to left <= series <= right.
# Series.clip([lower, upper, axis, inplace])	Trim values at input threshold(s).
# Series.clip_lower(threshold[, axis, inplace])	Return copy of the input with values below given value(s) truncated.
# Series.clip_upper(threshold[, axis, inplace])	Return copy of input with values above given value(s) truncated.
# Series.corr(other[, method, min_periods])	Compute correlation with other Series, excluding missing values
# Series.count([level])	Return number of non-NA/null observations in the Series
# Series.cov(other[, min_periods])	Compute covariance with Series, excluding missing values
# Series.cummax([axis, skipna])	Return cumulative max over requested axis.
# Series.cummin([axis, skipna])	Return cumulative minimum over requested axis.
# Series.cumprod([axis, skipna])	Return cumulative product over requested axis.
# Series.cumsum([axis, skipna])	Return cumulative sum over requested axis.
# Series.describe([percentiles, include, exclude])	Generates descriptive statistics that summarize the central tendency, dispersion and shape of a datasets distribution, excluding NaN values.
# Series.diff([periods])	1st discrete difference of object
# Series.factorize([sort, na_sentinel])	Encode the object as an enumerated type or categorical variable
# Series.kurt([axis, skipna, level, numeric_only])	Return unbiased kurtosis over requested axis using Fisher�s definition of kurtosis (kurtosis of normal == 0.0).
# Series.mad([axis, skipna, level])	Return the mean absolute deviation of the values for the requested axis
# Series.max([axis, skipna, level, numeric_only])	This method returns the maximum of the values in the object.
# Series.mean([axis, skipna, level, numeric_only])	Return the mean of the values for the requested axis
# Series.median([axis, skipna, level, ...])	Return the median of the values for the requested axis
# Series.min([axis, skipna, level, numeric_only])	This method returns the minimum of the values in the object.
# Series.mode()	Return the mode(s) of the dataset.
# Series.nlargest([n, keep])	Return the largest n elements.
# Series.nsmallest([n, keep])	Return the smallest n elements.
# Series.pct_change([periods, fill_method, ...])	Percent change over given number of periods.
# Series.prod([axis, skipna, level, ...])	Return the product of the values for the requested axis
# Series.quantile([q, interpolation])	Return value at the given quantile, a la numpy.percentile.
# Series.rank([axis, method, numeric_only, ...])	Compute numerical data ranks (1 through n) along axis.
# Series.sem([axis, skipna, level, ddof, ...])	Return unbiased standard error of the mean over requested axis.
# Series.skew([axis, skipna, level, numeric_only])	Return unbiased skew over requested axis
# Series.std([axis, skipna, level, ddof, ...])	Return sample standard deviation over requested axis.
# Series.sum([axis, skipna, level, ...])	Return the sum of the values for the requested axis
# Series.var([axis, skipna, level, ddof, ...])	Return unbiased variance over requested axis.
# Series.unique()	Return unique values in the object.
# Series.nunique([dropna])	Return number of unique elements in the object.
# Series.is_unique	Return boolean if values in the object are unique
# Series.is_monotonic	Return boolean if values in the object are
# Series.is_monotonic_increasing	Return boolean if values in the object are
# Series.is_monotonic_decreasing	Return boolean if values in the object are
# Series.value_counts([normalize, sort, ...])	Returns object containing counts of unique values.


# =====================================================================================================
# =====================================================================================================
class GroupBy(GroupByOps):
    """
    Parameters
    ----------
    dataset: Dataset
        The dataset object

    keys: list. List of column names to groupby

    filter: None. Boolean mask array applied as filter before grouping

    return_all: bool.  Default to False. When set to True will return all
          the dataset columns for every operation.

    hint_size: int.  Hint size for the hash (optional)

    sort_display: bool. Default to True.  Indicates

    lex: bool
        Defaults to False. When True uses a lexsort to find the groups (otherwise uses a hash).

    totals: bool
    """

    DebugMode = False
    TestCatGb = True

    def __init__(
        self,
        dataset,
        keys: list = None,
        filter=None,
        ordered=None,
        sort_display=None,
        return_all=False,
        hint_size=0,
        lex: bool = None,
        rec: bool = False,
        totals=False,
        copy=False,
        cutoffs=None,
        verbose=False,
        **kwargs,
    ):

        # upon creation create a unique id to track sorting
        self._uniqueid = GetTSC()

        # short circuit in copy mode
        if copy:
            return

        # copy over sort list
        self._col_sortlist = dataset._col_sortlist

        grouping_dict = None

        if not isinstance(keys, list):
            # if just one str passed in, wrap it in a list
            if isinstance(keys, str):
                keys = [keys]
            else:
                if isinstance(keys, np.ndarray):
                    if len(keys) != dataset.shape[0]:
                        raise TypeError(
                            f'The argument "keys" passed in GroupBy must be the same length as the dataset {len(keys)} vs {dataset.shape[0]}.'
                        )
                    else:
                        name = "col_0"
                        try:
                            name = keys.get_name()
                        except:
                            pass
                        # make into a dict
                        keys = {name: keys}
                        grouping_dict = keys
                else:
                    raise TypeError(
                        f'The argument "keys" passed in GroupBy init must be a list or str.  It must contain column names.'
                    )

        self._sortedbykey = False
        self._isortrows = None
        # self._gbkeys = None
        self._gb_keychain = None

        self._return_all = return_all
        self._filter = filter
        self._pcutoffs = cutoffs
        self._totals = totals

        if grouping_dict is None:
            grouping_dict = {k: dataset[k] for k in keys}

        if ordered is not None and sort_display is not None:
            raise TypeError(f"The kwarg ordered and sort_display cannot be used in combination")

        if ordered is not None:
            sort_display = ordered

        if sort_display is None:
            # automatically sort display if nothing passed
            if ordered is None and lex is None:
                sort_display = True
                try:
                    # if the first key has a grouping object and it does not want the display sorted and is itself not sorted,
                    # then we do not sort the display
                    # this happens after a 'cut' where we want to preserve label order
                    for v in grouping_dict.values():
                        if not v.grouping.isdisplaysorted and not v.grouping.isordered:
                            sort_display = False
                        break
                except:
                    pass

        self.grouping = TypeRegister.Grouping(
            grouping_dict,
            hint_size=hint_size,
            filter=filter,
            sort_display=sort_display,
            lex=lex,
            rec=rec,
            cutoffs=cutoffs,
            verbose=verbose,
        )

        # always use the names of the keys in the grouping dict
        # grouping dict should only live in grouping
        gvals = list(self.grouping._anydict.values())
        if len(keys) == len(gvals):
            self._grouping_dict = {k: gvals[i] for i, k in enumerate(keys)}
        # grouping dict may be larger than keys (multikey categorical splits into multiple arrays)
        else:
            self._grouping_dict = self.grouping._anydict

        # todo: remember change count on dataset
        self._dataset = dataset

        self._sort_display = sort_display

    # ---------------------------------------------------------------
    def copy(self, deep=True):
        """
        Called from getitem when user follows gb with []
        """
        newgb = type(self)(None, copy=True)
        newgb._col_sortlist = self._col_sortlist
        newgb._sortedbykey = self._sortedbykey
        try:
            newgb._gb_keychain = self._gb_keychain.copy(deep=deep)
        except:
            newgb._gb_keychain = None

        newgb._grouping_dict = self._grouping_dict.copy()
        newgb.grouping = self.grouping.copy(deep=deep)
        if deep:
            newgb._dataset = self._dataset.copy()
            if self._filter is None:
                newgb._filter = self._filter
            else:
                newgb._filter = self._filter.copy()

            # newgb.grouping = self.grouping.copy(grouping_dict=newgb._grouping_dict)

        else:
            newgb._dataset = self._dataset
            newgb._filter = self._filter

            # reference vs copy
            # newgb.grouping = self.grouping

        newgb._return_all = self._return_all
        newgb._pcutoffs = self._pcutoffs

        newgb._sort_display = self._sort_display
        # newgb._gbkeys = self._gbkeys
        newgb._isortrows = self._isortrows
        newgb._totals = self._totals

        return newgb

    # ---------------------------------------------------------------
    @property
    def gbkeys(self):
        """
        dictionary of numpy arrays binned from
        """
        return self.gb_keychain.gbkeys

    # ---------------------------------------------------------------
    @property
    def isortrows(self):
        """
        sorted index or None
        """
        return self.gb_keychain.isortrows

    # ---------------------------------------------------------------
    @property
    def ifirstkey(self):
        return self.grouping.ifirstkey

    # ---------------------------------------------------------------
    @property
    def ilastkey(self):
        return self.grouping.ilastkey

    # ---------------------------------------------------------------
    @property
    def gb_keychain(self):
        if self._gb_keychain is None:
            # this will change when groupby keys are handled only by the Grouping object
            # self._gb_keychain = GroupByKeys(self.grouping._anydict, self.ifirstkey, sort_display=self._sort_display, prebinned=self.grouping.iscategorical)
            # always send in the uniques
            self._gb_keychain = GroupByKeys(
                self.grouping.uniquedict, None, sort_display=self._sort_display, prebinned=True
            )
        return self._gb_keychain

    # ---------------------------------------------------------------
    def __iter__(self):
        """
        Generates tuples of key, value pairs.
        Keys are key values for single key, or tuples of key values for multikey.
        Values are datasets containing all rows from data in group for that key.
        """
        return self._iter_internal_contiguous()

    # ---------------------------------------------------------------
    def _pop_gb_data(self, calledfrom, userfunc, *args, **kwargs):
        """
        GroupBy holds on to its dataset. There may be no additional data provided.
        """
        kwargs.setdefault("dataset", self._dataset)
        return super()._pop_gb_data(calledfrom, userfunc, *args, **kwargs)

    # ---------------------------------------------------------------
    def add_totals(self, gb_ds):
        total = {}
        labels = gb_ds.label_get_names()
        for colname, arr in gb_ds.items():
            if colname not in labels:
                if arr.iscomputable():
                    total[colname] = arr.sum()
                else:
                    # leave blank
                    total[colname] = ""

        gb_ds.footer_set_values("Total", total)
        return gb_ds

    # ---------------------------------------------------------------
    # OVERRIDEN from groupbyops
    def _calculate_all(self, funcNum, *args, func_param=0, **kwargs):
        """
        Generate a GroupByKeys object if necessary and ask for the result of a calculation from the grouping object.
        Returns: a grouped by dataset with the result from the calculation
        """
        keychain = self.gb_keychain
        kwargs.setdefault("dataset", self._dataset)
        return_all = kwargs.pop("return_all", self._return_all)

        origdict, user_args, tups = self._prepare_gb_data("GroupBy", funcNum, *args, **kwargs)
        gb_ds = self.grouping._calculate_all(
            origdict,
            funcNum,
            func_param=func_param,
            return_all=return_all,
            keychain=keychain,
            user_args=user_args,
            tups=tups,
            **kwargs,
        )

        # check for accum1 like functionality
        if self._totals:
            return self.add_totals(gb_ds)
        else:
            return self._possibly_transform(gb_ds, label_keys=keychain.keys(), **kwargs)

    # ---------------------------------------------------------------
    def count(self, **kwargs):
        """Compute count of group"""
        return self.grouping.count(keychain=self.gb_keychain, **kwargs)

    ## ------------------------------------------------------------
    def __getattr__(self, name):
        """
        __getattr__ is hit when '.' is used to trim a single column.

        Examples
        --------
        >>> ds = Dataset({'col_'+str(i): np.random.rand(5) for i in range(5)})
        >>> ds.keycol = FA(['a','a','b','c','a'])
        >>> ds.gb('keycol').col_4.mean()
        *keycol   col_4
        -------   -----
        a          0.73
        b          0.03
        c          0.76

        """
        # check if the name exists in the dataset
        if name[0] == "_" or name not in self._dataset:
            return super().__getattribute__(name)

        return self._getitem(name)

    # ---------------------------------------------------------------
    def as_categorical(self):
        """
        Returns a categorical using the same binning information as the GroupBy object (no addtl. hash required).
        New categorical will not share a grouping object with this groupby object, but will share a reference to the iKey.
        Categorical operation results will be sorted or unsorted depending on if 'gb' or 'gbu' called this.
        """

        idx = self.grouping.ikey
        gbkeys = self.gbkeys

        return TypeRegister.Categorical(idx, _from_categorical=gbkeys, sort_gb=self._sort_display)

    # ---------------------------------------------------------------
    def _grouping_data_as_dict(self, ds):
        return self._dataset.as_ordered_dictionary()

    # -------------------------------------------------------
    def pad(self, limit=0, fill_val=None, inplace=False):
        """
        Forward fill the values

        Parameters
        ----------
        limit : integer, optional
            limit of how many values to fill

        See Also
        --------
        fill_forward
        fill_backward
        fill_invalid
        """
        return self.fill_forward(fill_val=fill_val, limit=limit, inplace=inplace)

    # -------------------------------------------------------
    def fill_forward(self, limit=0, fill_val=None, inplace=False):
        """
        Replace NaN and invalid array values by propagating the last encountered valid
        group value forward.

        Parameters
        ----------
        limit : int, default 0
            The maximium number of consecutive NaN or invalid values to fill. If there
            is a gap with more than this number of consecutive NaN or invalid values,
            the gap will be only partially filled. If no `limit` is specified, all
            consecutive NaN and invalid values are replaced.
        fill_val : scalar, default None
            The value to use where there is no valid group value to propagate forward.
            If `fill_val` is not specified, NaN and invalid values aren't replaced where
            there is no valid group value to propagate forward.
        inplace : bool, default False
            If False, return a copy of the array. If True, modify original data. This
            will modify any other views on this object. This fails if the array is
            locked.

        Returns
        -------
        `Dataset`
            The returned `Dataset` contains the input `Dataset` object's numerical
            columns.

        See Also
        --------
        GroupBy.fill_backward :
            Replace NaN and invalid array values with the next valid group value.
        Categorical.fill_forward :
            Replace NaN and invalid array values with the last valid group value.
        riptable.fill_forward : Replace NaN and invalid values with the last valid
            value.
        Dataset.fillna : Replace NaN and invalid values with a specified value or
            nearby data.
        FastArray.fillna : Replace NaN and invalid values with a specified value or
            nearby data.
        FastArray.replacena : Replace NaN and invalid values with a specified value.

        Examples
        --------
        >>> ds = rt.Dataset({'Key_col' : ['A', 'B', 'A', 'B', 'A', 'B'],
        ...                  'Vals' : [0, 1, 2, 3, rt.nan, rt.nan]})
        >>> ds.gb('Key_col').fill_forward()
        #   Vals
        -   ----
        0   0.00
        1   1.00
        2   2.00
        3   3.00
        4   2.00
        5   3.00

        Use a `fill_val` to replace values where there's no valid group value to
        propagate forward:

        >>> ds.Vals = rt.FastArray([rt.nan, rt.nan, 2, 3, 4, 5])
        >>> ds.gb('Key_col').fill_forward(fill_val = 0)
        #   Vals
        -   ----
        0   0.00
        1   0.00
        2   2.00
        3   3.00
        4   4.00
        5   5.00


        Replace only the first NaN or invalid value in any consecutive series of NaN or
        invalid values in a group:

        >>> ds.Vals = rt.FastArray([0, 1, rt.nan, rt.nan, rt.nan, rt.nan])
        >>> ds.gb('Key_col').fill_forward(limit = 1)
        #   Vals
        -   ----
        0   0.00
        1   1.00
        2   0.00
        3   1.00
        4    nan
        5    nan
        """
        return super().nb_fill_forward(fill_val=fill_val, limit=limit, inplace=inplace)

    # -------------------------------------------------------
    def backfill(self, limit=0, fill_val=None, inplace=False):
        """
        Backward fill the values

        Parameters
        ----------
        limit : integer, optional
            limit of how many values to fill

        See Also
        --------
        fill_forward
        fill_backward
        fill_invalid
        """
        # stub for pandas
        return self.fill_backward(fill_val=fill_val, limit=limit, inplace=inplace)

    # -------------------------------------------------------
    def fill_backward(self, limit=0, fill_val=None, inplace=False):
        """
        Replace NaN and invalid array values by propagating the next encountered valid
        group value backward.

        Parameters
        ----------
        limit : int, default 0
            The maximium number of consecutive NaN or invalid values to fill. If there
            is a gap with more than this number of consecutive NaN or invalid values,
            the gap will be only partially filled. If no `limit` is specified, all
            consecutive NaN and invalid values are replaced.
        fill_val : scalar, default None
            The value to use where there is no valid group value to propagate backward.
            If `fill_val` is not specified, NaN and invalid values aren't replaced where
            there is no valid group value to propagate backward.
        **kwargs
            Additional keyword arguments.

        Returns
        -------
        `Dataset`
            The returned `Dataset` contains the input `Dataset` object's numerical
            columns.

        See Also
        --------
        GroupBy.fill_forward :
            Replace NaN and invalid array values with the last valid group value.
        Categorical.fill_backward :
            Replace NaN and invalid array values with the next valid group value.
        riptable.fill_backward : Replace NaN and invalid values with the next valid
            value.
        Dataset.fillna : Replace NaN and invalid values with a specified value or
            nearby data.
        FastArray.fillna : Replace NaN and invalid values with a specified value or
            nearby data.
        FastArray.replacena : Replace NaN and invalid values with a specified value.

        Examples
        --------
        >>> ds = rt.Dataset({'Key_col' : ['A', 'B', 'A', 'B', 'A', 'B'],
        ...             'Vals' : [rt.nan, rt.nan, 2, 3, 4, 5]})
        >>> ds.gb('Key_col').fill_backward()
        #   Vals
        -   ----
        0   2.00
        1   3.00
        2   2.00
        3   3.00
        4   4.00
        5   5.00

        Use a `fill_val` to replace values where there's no valid group value to
        propagate backward:

        >>> ds.Vals = rt.FastArray([0, 1, 2, 3, rt.nan, rt.nan])
        >>> ds.gb('Key_col').fill_backward(fill_val = 0)
        #   Vals
        -   ----
        0   0.00
        1   1.00
        2   2.00
        3   3.00
        4   0.00
        5   0.00

        Replace only the first NaN or invalid value in any consecutive series of NaN or
        invalid values in a group:

        >>> ds.Vals = rt.FastArray([rt.nan, rt.nan, rt.nan, rt.nan, 4, 5])
        >>> ds.gb('Key_col').fill_backward(limit = 1)
        #   Vals
        -   ----
        0    nan
        1    nan
        2   4.00
        3   5.00
        4   4.00
        5   5.00
        """
        return super().nb_fill_backward(fill_val=fill_val, limit=limit, inplace=inplace)

    # -------------------------------------------------------
    def expanding(self, **kwargs):
        raise NotImplementedError("riptable uses the Grouping object (see .group)ck")

    # -------------------------------------------------------
    def stack(self, **kwargs):
        raise NotImplementedError("riptable uses the Accum2 class for stack/unstack")

    # -------------------------------------------------------
    def unstack(self, **kwargs):
        raise NotImplementedError("riptable uses the Accum2 class for stack/unstack")

    # -------------------------------------------------------
    @property
    def transform(self):
        """
        The property transform sets a flag so that the next reduce function called after transform,
        will repopulate the original array with the reduced value.

        Example
        -------
        >>> ds.groupby(['side', 'venue']).transform.sum()
        """
        warnings.warn("Deprecation warning: Use kwarg transform=True instead of transform.")
        self._transform = True
        return self

    # -------------------------------------------------------
    def get_group(self, category, **kwargs):
        """
        The name of the group to get as a Dataset.

        Parameters
        ----------
        category: string or tuple
            A value from the column used to construct the GroupBy, or if
            multiple columns were used, a tuple of the multiple columns.

        Returns
        -------
        Dataset

        Example
        -------
        >>> ds.groupby('symbol').get_group('AAPL')
        """
        # categorical has method to get the bin
        cat = self.as_categorical()

        bin = cat.from_category(category)
        return self._dataset[self.grouping.ikey == bin, :]

    # -------------------------------------------------------
    def _build_string(self):
        newString = f"GroupBy Keys {list(self._grouping_dict.keys())} @ [{self._dataset._ncols} x {self.grouping.unique_count}]\n"
        newString += f"ikey:{not isinstance(self.grouping.ikey, type(None))}  iFirstKey:{not isinstance(self.grouping.iFirstKey, type(None))}  iNextKey:{not isinstance(self.grouping.iNextKey, type(None))}  nCountGroup:{not isinstance(self.grouping.nCountGroup, type(None))} "
        # newString += f" _isortrows:{not isinstance(self._isortrows, type(None))}  _gbkeys:{not isinstance(self._gbkeys, type(None))}  _filter:{not isinstance(self._filter, type(None))}  _return_all:{self._return_all}\n\n"

        newString += f"_filter:{not isinstance(self._filter, type(None))}  _return_all:{self._return_all}\n\n"

        # temporarily build a dataset
        dset = self.count()

        # dump of count
        resultString = dset.__str__()

        return newString + resultString

    # -------------------------------------------------------
    def __repr__(self):
        return self._build_string()

    # -------------------------------------------------------
    def __str__(self):
        return self._build_string()

    # -------------------------------------------------------
    def _getitem(self, fld):
        """
        Called by __getitem__ and __getattr__.
        Uses the field to index into the stored dataset.
        Often used to limit the data the groupby operation is being performed on.
        Returns a shallow copy of the groupby object.

        This routine gets hit during the following common code pattern:

        >>> ds = Dataset({'col_'+str(i): np.random.rand(5) for i in range(5)})
        >>> ds.keycol = FA(['a','a','b','c','a'])
        >>> ds.gb('keycol')[['col_1', 'col_2']].sum()
        *keycol   col_1   col_2
        -------   -----   -----
        a          1.92    0.89
        b          0.70    0.46
        c          0.07    0.42

        >>> ds.gb('keycol').col_4.mean()
        *keycol   col_4
        -------   -----
        a          0.73
        b          0.03
        c          0.76

        """
        # return new groupby
        newgb = self.copy(deep=False)

        # reduce the datasets
        if isinstance(fld, (str, bytes)):
            # need a dataset back, not a fastarray
            index = [fld]
        elif isinstance(fld, list):
            index = fld.copy()
        # turn to list of multiple strings
        elif isinstance(fld, tuple):
            index = list(fld)
        else:
            raise TypeError("Dont know how to make groupby from", fld)

        # we must add back in the grouping keys or we cannot group
        for item in self._grouping_dict.keys():
            if item not in index:
                index.append(item)

        newgb._dataset = newgb._dataset[index]
        return newgb

    # -------------------------------------------------------
    def __getitem__(self, fld):
        return self._getitem(fld)


# keep this as the last line
from .rt_enum import TypeRegister

TypeRegister.GroupBy = GroupBy
