# Copyright (c) 2020, NVIDIA CORPORATION
import itertools

import cupy as cp
import numba
import pandas as pd

import cudf
from cudf import _lib as libcudf
from cudf.core.column.column import as_column
from cudf.utils import cudautils


class Rolling:
    """
    Rolling window calculations.

    Parameters
    ----------
    window : int or offset
        Size of the window, i.e., the number of observations used
        to calculate the statistic.
        For datetime indexes, an offset can be provided instead
        of an int. The offset must be convertible to a timedelta.
        As opposed to a fixed window size, each window will be
        sized to accommodate observations within the time period
        specified by the offset.
    min_periods : int, optional
        The minimum number of observations in the window that are
        required to be non-null, so that the result is non-null.
        If not provided or ``None``, ``min_periods`` is equal to
        the window size.
    center : bool, optional
        If ``True``, the result is set at the center of the window.
        If ``False`` (default), the result is set at the right edge
        of the window.

    Returns
    -------
    ``Rolling`` object.

    Examples
    --------
    >>> import cudf
    >>> a = cudf.Series([1, 2, 3, None, 4])

    Rolling sum with window size 2.

    >>> print(a.rolling(2).sum())
    0
    1    3
    2    5
    3
    4
    dtype: int64

    Rolling sum with window size 2 and min_periods 1.

    >>> print(a.rolling(2, min_periods=1).sum())
    0    1
    1    3
    2    5
    3    3
    4    4
    dtype: int64

    Rolling count with window size 3.

    >>> print(a.rolling(3).count())
    0    1
    1    2
    2    3
    3    2
    4    2
    dtype: int64

    Rolling count with window size 3, but with the result set at the
    center of the window.

    >>> print(a.rolling(3, center=True).count())
    0    2
    1    3
    2    2
    3    2
    4    1 dtype: int64

    Rolling max with variable window size specified by an offset;
    only valid for datetime index.

    >>> a = cudf.Series(
    ...     [1, 9, 5, 4, np.nan, 1],
    ...     index=[
    ...         pd.Timestamp('20190101 09:00:00'),
    ...         pd.Timestamp('20190101 09:00:01'),
    ...         pd.Timestamp('20190101 09:00:02'),
    ...         pd.Timestamp('20190101 09:00:04'),
    ...         pd.Timestamp('20190101 09:00:07'),
    ...         pd.Timestamp('20190101 09:00:08')
    ...     ]
    ... )

    >>> print(a.rolling('2s').max())
    2019-01-01T09:00:00.000    1
    2019-01-01T09:00:01.000    9
    2019-01-01T09:00:02.000    9
    2019-01-01T09:00:04.000    4
    2019-01-01T09:00:07.000
    2019-01-01T09:00:08.000    1
    dtype: int64

    Apply custom function on the window with the *apply* method

    >>> import numpy as np
    >>> import math
    >>> b = cudf.Series([16, 25, 36, 49, 64, 81], dtype=np.float64)
    >>> def some_func(A):
    ...     b = 0
    ...     for a in A:
    ...         b = b + math.sqrt(a)
    ...     return b
    ...
    >>> print(b.rolling(3, min_periods=1).apply(some_func))
    0     4.0
    1     9.0
    2    15.0
    3    18.0
    4    21.0
    5    24.0
    dtype: float64

    And this also works for window rolling set by an offset

    >>> import pandas as pd
    >>> c = cudf.Series(
    ...     [16, 25, 36, 49, 64, 81],
    ...     index=[
    ...          pd.Timestamp('20190101 09:00:00'),
    ...          pd.Timestamp('20190101 09:00:01'),
    ...          pd.Timestamp('20190101 09:00:02'),
    ...          pd.Timestamp('20190101 09:00:04'),
    ...          pd.Timestamp('20190101 09:00:07'),
    ...          pd.Timestamp('20190101 09:00:08')
    ...      ],
    ...     dtype=np.float64
    ... )
    >>> print(c.rolling('2s').apply(some_func))
    2019-01-01T09:00:00.000     4.0
    2019-01-01T09:00:01.000     9.0
    2019-01-01T09:00:02.000    11.0
    2019-01-01T09:00:04.000     7.0
    2019-01-01T09:00:07.000     8.0
    2019-01-01T09:00:08.000    17.0
    dtype: float64
    """

    _time_window = False

    def __init__(
        self,
        obj,
        window,
        min_periods=None,
        center=False,
        axis=0,
        win_type=None,
    ):
        self.obj = obj
        self.window = window
        self.min_periods = min_periods
        self.center = center
        self._normalize()
        if axis != 0:
            raise NotImplementedError("axis != 0 is not supported yet.")
        self.axis = axis

        if win_type is not None:
            if win_type != "boxcar":
                raise NotImplementedError(
                    "Only the default win_type 'boxcar' is currently supported"
                )
        self.win_type = win_type

    def __getattr__(self, key):
        if key == "obj":
            raise AttributeError()
        return self.obj[key].rolling(
            window=self.window,
            min_periods=self.min_periods,
            center=self.center,
        )

    def __getitem__(self, arg):
        if isinstance(arg, tuple):
            arg = list(arg)
        return self.obj[arg].rolling(
            window=self.window,
            min_periods=self.min_periods,
            center=self.center,
        )

    def _apply_agg_series(self, sr, agg_name):
        if isinstance(self.window, int):
            result_col = libcudf.rolling.rolling(
                sr._column,
                None,
                None,
                self.window,
                self.min_periods,
                self.center,
                agg_name,
            )
        else:
            from cudf.utils import cudautils

            result_col = libcudf.rolling.rolling(
                sr._column,
                as_column(self.window),
                as_column(
                    cudautils.full(self.window.size, 0, self.window.dtype)
                ),
                None,
                self.min_periods,
                self.center,
                agg_name,
            )
        return sr._copy_construct(data=result_col)

    def _apply_agg_dataframe(self, df, agg_name):
        result_df = cudf.DataFrame({})
        for i, col_name in enumerate(df.columns):
            result_col = self._apply_agg_series(df[col_name], agg_name)
            result_df.insert(i, col_name, result_col)
        result_df.index = df.index
        return result_df

    def _apply_agg(self, agg_name):
        if agg_name == "count" and not self._time_window:
            self.min_periods = 0
        if isinstance(self.obj, cudf.Series):
            return self._apply_agg_series(self.obj, agg_name)
        else:
            return self._apply_agg_dataframe(self.obj, agg_name)

    def sum(self):
        return self._apply_agg("sum")

    def min(self):
        return self._apply_agg("min")

    def max(self):
        return self._apply_agg("max")

    def mean(self):
        return self._apply_agg("mean")

    def count(self):
        return self._apply_agg("count")

    def apply(self, func, *args, **kwargs):
        """
        Counterpart of pandas.core.window.Rolling.apply

        *func* is a user defined function that takes an 1D array as input:

        See also
        --------
        The Notes section in `Series.applymap`.

        """
        has_nulls = False
        if isinstance(self.obj, cudf.Series):
            if self.obj._column.has_nulls:
                has_nulls = True
        else:
            for col in self.obj._data:
                if self.obj[col].has_nulls:
                    has_nulls = True
        if has_nulls:
            raise NotImplementedError(
                "Handling UDF with null values is not yet supported"
            )
        return self._apply_agg(func)

    def _normalize(self):
        """
        Normalize the *window* and *min_periods* args

        *window* can be:

        * An integer, in which case it is the window size.
          If *min_periods* is unspecified, it is set to be equal to
          the window size.

        * A timedelta offset, in which case it is used to generate
          a column of window sizes to use for each element.
          If *min_periods* is unspecified, it is set to 1.
          Only valid for datetime index.
        """
        window, min_periods = self.window, self.min_periods
        if pd.api.types.is_number(window):
            # only allow integers
            if not pd.api.types.is_integer(window):
                raise ValueError("window must be an integer")
            if window <= 0:
                raise ValueError("window cannot be zero or negative")
            if self.min_periods is None:
                min_periods = window
        else:
            if isinstance(window, numba.cuda.devicearray.DeviceNDArray):
                # window is a device_array of window sizes
                self.window = window
                self.min_periods = min_periods
                return

            if not isinstance(self.obj.index, cudf.core.index.DatetimeIndex):
                raise ValueError(
                    "window must be an integer for " "non datetime index"
                )

            self._time_window = True

            try:
                window = pd.to_timedelta(window)
                # to_timedelta will also convert np.arrays etc.,
                if not isinstance(window, pd.Timedelta):
                    raise ValueError
                window = window.to_timedelta64()
            except ValueError as e:
                raise ValueError(
                    "window must be integer or " "convertible to a timedelta"
                ) from e
            if self.min_periods is None:
                min_periods = 1

        self.window = self._window_to_window_sizes(window)
        self.min_periods = min_periods

    def _window_to_window_sizes(self, window):
        """
        For non-fixed width windows,
        convert the window argument into window sizes.
        """
        if pd.api.types.is_integer(window):
            return window
        else:
            return cudautils.window_sizes_from_offset(
                self.obj.index._values.data_array_view, window
            )

    def __repr__(self):
        return "{} [window={},min_periods={},center={}]".format(
            self.__class__.__name__, self.window, self.min_periods, self.center
        )


class RollingGroupby(Rolling):
    def __init__(self, groupby, window, min_periods=None, center=False):
        """
        Grouped rolling window calculation.

        See also
        --------
        cudf.core.window.Rolling
        """
        sort_order = groupby.grouping.keys.argsort()

        # TODO: there may be overlap between the columns
        # of `groupby.grouping.keys` and `groupby.obj`.
        # As an optimization, avoid gathering those twice.
        self._group_keys = groupby.grouping.keys.take(sort_order)
        obj = groupby.obj.take(sort_order)

        gb_size = groupby.size()
        self._group_starts = (
            gb_size.cumsum().shift(1).fillna(0).repeat(gb_size)
        )

        super().__init__(obj, window, min_periods=min_periods, center=center)

    def _window_to_window_sizes(self, window):
        if pd.api.types.is_integer(window):
            return cudautils.grouped_window_sizes_from_offset(
                cp.arange(len(self.obj)), self._group_starts, window
            )
        else:
            return cudautils.grouped_window_sizes_from_offset(
                self.obj.index._values.data_array_view,
                self._group_starts,
                window,
            )

    def _apply_agg(self, agg_name):
        index = cudf.MultiIndex.from_frame(
            cudf.DataFrame(
                {
                    key: value
                    for key, value in itertools.chain(
                        self._group_keys._data.items(),
                        self.obj.index._data.items(),
                    )
                }
            )
        )

        result = super()._apply_agg(agg_name).set_index(index)
        return result
