# Source code for dask.dataframe.rolling

    from __future__ import annotations

    import datetime
    import warnings
    from numbers import Integral

    import pandas as pd
    from pandas.api.types import is_datetime64_any_dtype
    from pandas.core.window import Rolling as pd_Rolling

    from dask.array.core import normalize_arg
    from dask.base import tokenize
    from dask.blockwise import BlockwiseDepDict, _blockwise_unpack_collections_task_spec
    from dask.dataframe import methods
    from dask.dataframe._compat import check_axis_keyword_deprecation
    from dask.dataframe.core import (
        Scalar,
        _Frame,
        _get_divisions_map_partitions,
        _get_meta_map_partitions,
        _maybe_from_pandas,
        apply_and_enforce,
        new_dd_object,
        partitionwise_graph,
    )
    from dask.dataframe.io import from_pandas
    from dask.dataframe.multi import _maybe_align_partitions
    from dask.dataframe.utils import (
        insert_meta_param_description,
        is_dask_collection,
        is_dataframe_like,
        is_series_like,
    )
    from dask.highlevelgraph import HighLevelGraph
    from dask.typing import no_default
    from dask.utils import M, apply, derived_from, funcname, has_keyword

    CombinedOutput = type("CombinedOutput", (tuple,), {})

    def _combined_parts(prev_part, current_part, next_part, before, after):
        msg = (
            "Partition size is less than overlapping "
            "window size. Try using ``df.repartition`` "
            "to increase the partition size."
        )

        if prev_part is not None and isinstance(before, Integral):
            if prev_part.shape[0] != before:
                raise NotImplementedError(msg)

        if next_part is not None and isinstance(after, Integral):
            if next_part.shape[0] != after:
                raise NotImplementedError(msg)

        parts = [p for p in (prev_part, current_part, next_part) if p is not None]
        combined = methods.concat(parts)

        return CombinedOutput(
            (
                combined,
                len(prev_part) if prev_part is not None else None,
                len(next_part) if next_part is not None else None,
            )
        )

    def overlap_chunk(func, before, after, *args, **kwargs):
        dfs = [df for df in args if isinstance(df, CombinedOutput)]
        combined, prev_part_length, next_part_length = dfs[0]

        args = [arg[0] if isinstance(arg, CombinedOutput) else arg for arg in args]

        out = func(*args, **kwargs)

        if prev_part_length is None:
            before = None
        if isinstance(before, datetime.timedelta):
            before = prev_part_length

        expansion = None
        if combined.shape[0] != 0:
            expansion = out.shape[0] // combined.shape[0]
        if before and expansion:
            before *= expansion
        if next_part_length is None:
            return out.iloc[before:]
        if isinstance(after, datetime.timedelta):
            after = next_part_length
        if after and expansion:
            after *= expansion
        return out.iloc[before:-after]

    @insert_meta_param_description
    def map_overlap(
        func,
        df,
        before,
        after,
        *args,
        meta=no_default,
        enforce_metadata=True,
        transform_divisions=True,
        align_dataframes=True,
        **kwargs,
    ):
        """Apply a function to each partition, sharing rows with adjacent partitions.

        Parameters
        ----------
        func : function
            The function applied to each partition. If this function accepts
            the special ``partition_info`` keyword argument, it will receive
            information on the partition's relative location within the
            dataframe.
        df: dd.DataFrame, dd.Series
        args, kwargs :
            Positional and keyword arguments to pass to the function.
            Positional arguments are computed on a per-partition basis, while
            keyword arguments are shared across all partitions. The partition
            itself will be the first positional argument, with all other
            arguments passed *after*. Arguments can be ``Scalar``, ``Delayed``,
            or regular Python objects. DataFrame-like args (both dask and
            pandas) will be repartitioned to align (if necessary) before
            applying the function; see ``align_dataframes`` to control this
            behavior.
        enforce_metadata : bool, default True
            Whether to enforce at runtime that the structure of the DataFrame
            produced by ``func`` actually matches the structure of ``meta``.
            This will rename and reorder columns for each partition,
            and will raise an error if this doesn't work,
            but it won't raise if dtypes don't match.
        before : int, timedelta or string timedelta
            The rows to prepend to partition ``i`` from the end of
            partition ``i - 1``.
        after : int, timedelta or string timedelta
            The rows to append to partition ``i`` from the beginning
            of partition ``i + 1``.
        transform_divisions : bool, default True
            Whether to apply the function onto the divisions and apply those
            transformed divisions to the output.
        align_dataframes : bool, default True
            Whether to repartition DataFrame- or Series-like args
            (both dask and pandas) so their divisions align before applying
            the function. This requires all inputs to have known divisions.
            Single-partition inputs will be split into multiple partitions.

            If False, all inputs must have either the same number of partitions
            or a single partition. Single-partition inputs will be broadcast to
            every partition of multi-partition inputs.
        $META

        See Also
        --------
        dd.DataFrame.map_overlap
        """

        df = (
            from_pandas(df, 1)
            if (is_series_like(df) or is_dataframe_like(df)) and not is_dask_collection(df)
            else df
        )

        args = (df,) + args

        if isinstance(before, str):
            before = pd.to_timedelta(before)
        if isinstance(after, str):
            after = pd.to_timedelta(after)

        if isinstance(before, datetime.timedelta) or isinstance(after, datetime.timedelta):
            if not is_datetime64_any_dtype(df.index._meta_nonempty.inferred_type):
                raise TypeError(
                    "Must have a `DatetimeIndex` when using string offset "
                    "for `before` and `after`"
                )
        else:
            if not (
                isinstance(before, Integral)
                and before >= 0
                and isinstance(after, Integral)
                and after >= 0
            ):
                raise ValueError("before and after must be positive integers")

        name = kwargs.pop("token", None)
        parent_meta = kwargs.pop("parent_meta", None)

        assert callable(func)
        if name is not None:
            token = tokenize(meta, before, after, *args, **kwargs)
        else:
            name = "overlap-" + funcname(func)
            token = tokenize(func, meta, before, after, *args, **kwargs)
        name = f"{name}-{token}"

        if align_dataframes:
            args = _maybe_from_pandas(args)
            try:
                args = _maybe_align_partitions(args)
            except ValueError as e:
                raise ValueError(
                    f"{e}. If you don't want the partitions to be aligned, and are "
                    "calling `map_overlap` directly, pass `align_dataframes=False`."
                ) from e

        dfs = [df for df in args if isinstance(df, _Frame)]

        meta = _get_meta_map_partitions(args, dfs, func, kwargs, meta, parent_meta)

        if all(isinstance(arg, Scalar) for arg in args):
            layer = {
                (name, 0): (
                    apply,
                    func,
                    (tuple, [(arg._name, 0) for arg in args]),
                    kwargs,
                )
            }
            graph = HighLevelGraph.from_collections(name, layer, dependencies=args)
            return Scalar(graph, name, meta)

        args2 = []
        dependencies = []

        divisions = _get_divisions_map_partitions(
            align_dataframes, transform_divisions, dfs, func, args, kwargs
        )

        def _handle_frame_argument(arg):
            dsk = {}

            prevs_parts_dsk, prevs = _get_previous_partitions(arg, before)
            dsk.update(prevs_parts_dsk)

            nexts_parts_dsk, nexts = _get_nexts_partitions(arg, after)
            dsk.update(nexts_parts_dsk)

            name_a = "overlap-concat-" + tokenize(arg)
            for i, (prev, current, next) in enumerate(
                zip(prevs, arg.__dask_keys__(), nexts)
            ):
                key = (name_a, i)
                dsk[key] = (_combined_parts, prev, current, next, before, after)

            graph = HighLevelGraph.from_collections(name_a, dsk, dependencies=[arg])
            return new_dd_object(graph, name_a, meta, divisions)

        for arg in args:
            if isinstance(arg, _Frame):
                arg = _handle_frame_argument(arg)
                args2.append(arg)
                dependencies.append(arg)
                continue
            arg = normalize_arg(arg)
            arg2, collections = _blockwise_unpack_collections_task_spec(arg)
            if collections:
                args2.append(arg2)
                dependencies.extend(collections)
            else:
                args2.append(arg)

        kwargs3 = {}
        simple = True
        for k, v in kwargs.items():
            v = normalize_arg(v)
            v, collections = _blockwise_unpack_collections_task_spec(v)
            dependencies.extend(collections)
            kwargs3[k] = v
            if collections:
                simple = False

        if has_keyword(func, "partition_info"):
            partition_info = {
                (i,): {"number": i, "division": division}
                for i, division in enumerate(divisions[:-1])
            }

            args2.insert(0, BlockwiseDepDict(partition_info))
            orig_func = func

            def func(partition_info, *args, **kwargs):
                return orig_func(*args, **kwargs, partition_info=partition_info)

        if enforce_metadata:
            dsk = partitionwise_graph(
                apply_and_enforce,
                name,
                func,
                before,
                after,
                *args2,
                dependencies=dependencies,
                _func=overlap_chunk,
                _meta=meta,
                **kwargs3,
            )
        else:
            kwargs4 = kwargs if simple else kwargs3
            dsk = partitionwise_graph(
                overlap_chunk,
                name,
                func,
                before,
                after,
                *args2,
                **kwargs4,
                dependencies=dependencies,
            )

        graph = HighLevelGraph.from_collections(name, dsk, dependencies=dependencies)
        return new_dd_object(graph, name, meta, divisions)

    def _get_nexts_partitions(df, after):
        """
        Helper to get the nexts partitions required for the overlap
        """
        dsk = {}
        df_name = df._name

        timedelta_partition_message = (
            "Partition size is less than specified window. "
            "Try using ``df.repartition`` to increase the partition size"
        )

        name_b = "overlap-append-" + tokenize(df, after)
        if after and isinstance(after, Integral):
            nexts = []
            for i in range(1, df.npartitions):
                key = (name_b, i)
                dsk[key] = (M.head, (df_name, i), after)
                nexts.append(key)
            nexts.append(None)
        elif isinstance(after, datetime.timedelta):
            # TODO: Do we have a use-case for this? Pandas doesn't allow negative rolling windows
            deltas = pd.Series(df.divisions).diff().iloc[1:-1]
            if (after > deltas).any():
                raise ValueError(timedelta_partition_message)

            nexts = []
            for i in range(1, df.npartitions):
                key = (name_b, i)
                dsk[key] = (_head_timedelta, (df_name, i - 0), (df_name, i), after)
                nexts.append(key)
            nexts.append(None)
        else:
            nexts = [None] * df.npartitions
        return dsk, nexts

    def _get_previous_partitions(df, before):
        """
        Helper to get the previous partitions required for the overlap
        """
        dsk = {}
        df_name = df._name

        name_a = "overlap-prepend-" + tokenize(df, before)
        if before and isinstance(before, Integral):
            prevs = [None]
            for i in range(df.npartitions - 1):
                key = (name_a, i)
                dsk[key] = (M.tail, (df_name, i), before)
                prevs.append(key)

        elif isinstance(before, datetime.timedelta):
            # Assumes monotonic (increasing?) index
            divs = pd.Series(df.divisions)
            deltas = divs.diff().iloc[1:-1]

            # In the first case window-size is larger than at least one partition, thus it is
            # necessary to calculate how many partitions must be used for each rolling task.
            # Otherwise, these calculations can be skipped (faster)

            if (before > deltas).any():
                pt_z = divs[0]
                prevs = [None]
                for i in range(df.npartitions - 1):
                    # Select all indexes of relevant partitions between the current partition and
                    # the partition with the highest division outside the rolling window (before)
                    pt_i = divs[i + 1]

                    # lower-bound the search to the first division
                    lb = max(pt_i - before, pt_z)

                    first, j = divs[i], i
                    while first > lb and j > 0:
                        first = first - deltas[j]
                        j = j - 1

                    key = (name_a, i)
                    dsk[key] = (
                        _tail_timedelta,
                        [(df_name, k) for k in range(j, i + 1)],
                        (df_name, i + 1),
                        before,
                    )
                    prevs.append(key)

            else:
                prevs = [None]
                for i in range(df.npartitions - 1):
                    key = (name_a, i)
                    dsk[key] = (
                        _tail_timedelta,
                        [(df_name, i)],
                        (df_name, i + 1),
                        before,
                    )
                    prevs.append(key)
        else:
            prevs = [None] * df.npartitions
        return dsk, prevs

    def _head_timedelta(current, next_, after):
        """Return rows of ``next_`` whose index is before the last
        observation in ``current`` + ``after``.

        Parameters
        ----------
        current : DataFrame
        next_ : DataFrame
        after : timedelta

        Returns
        -------
        overlapped : DataFrame
        """
        return next_[next_.index < (current.index.max() + after)]

    def _tail_timedelta(prevs, current, before):
        """Return the concatenated rows of each dataframe in ``prevs`` whose
        index is after the first observation in ``current`` - ``before``.

        Parameters
        ----------
        current : DataFrame
        prevs : list of DataFrame objects
        before : timedelta

        Returns
        -------
        overlapped : DataFrame
        """
        selected = methods.concat(
            [prev[prev.index > (current.index.min() - before)] for prev in prevs]
        )
        return selected

    class Rolling:
        """Provides rolling window calculations."""

        def __init__(
            self,
            obj,
            window=None,
            min_periods=None,
            center=False,
            win_type=None,
            axis=no_default,
        ):
            self.obj = obj  # dataframe or series
            self.window = window
            self.min_periods = min_periods
            self.center = center
            self.axis = axis
            self.win_type = win_type
            # Allow pandas to raise if appropriate
            obj._meta.rolling(**self._rolling_kwargs())
            # Using .rolling(window='2s'), pandas will convert the
            # offset str to a window in nanoseconds. But pandas doesn't
            # accept the integer window with win_type='freq', so we store
            # that information here.
            # See https://github.com/pandas-dev/pandas/issues/15969
            self._win_type = None if isinstance(self.window, int) else "freq"

            if self.axis in ("index", 1, "rows"):
                warnings.warn(
                    "Using axis=1 in Rolling has been deprecated and will be removed "
                    "in a future version.",
                    FutureWarning,
                )

        def _rolling_kwargs(self):
            kwargs = {
                "window": self.window,
                "min_periods": self.min_periods,
                "center": self.center,
                "win_type": self.win_type,
            }
            if self.axis is not no_default:
                kwargs["axis"] = self.axis
            return kwargs

        @property
        def _has_single_partition(self):
            """
            Indicator for whether the object has a single partition (True)
            or multiple (False).
            """
            return (
                self.axis in (1, "columns")
                or (isinstance(self.window, Integral) and self.window <= 1)
                or self.obj.npartitions == 1
            )

        @staticmethod
        def pandas_rolling_method(df, rolling_kwargs, name, *args, **kwargs):
            with check_axis_keyword_deprecation():
                rolling = df.rolling(**rolling_kwargs)
            return getattr(rolling, name)(*args, **kwargs)

        def _call_method(self, method_name, *args, **kwargs):
            rolling_kwargs = self._rolling_kwargs()
            meta = self.pandas_rolling_method(
                self.obj._meta_nonempty, rolling_kwargs, method_name, *args, **kwargs
            )

            if self._has_single_partition:
                # There's no overlap just use map_partitions
                return self.obj.map_partitions(
                    self.pandas_rolling_method,
                    rolling_kwargs,
                    method_name,
                    *args,
                    token=method_name,
                    meta=meta,
                    **kwargs,
                )
            # Convert window to overlap
            if self.center:
                before = self.window // 2
                after = self.window - before - 1
            elif self._win_type == "freq":
                before = pd.Timedelta(self.window)
                after = 0
            else:
                before = self.window - 1
                after = 0
            return map_overlap(
                self.pandas_rolling_method,
                self.obj,
                before,
                after,
                rolling_kwargs,
                method_name,
                *args,
                token=method_name,
                meta=meta,
                **kwargs,
            )

        @derived_from(pd_Rolling)
        def count(self):
            return self._call_method("count")

        @derived_from(pd_Rolling)
        def cov(self):
            return self._call_method("cov")

        @derived_from(pd_Rolling)
        def sum(self):
            return self._call_method("sum")

        @derived_from(pd_Rolling)
        def mean(self):
            return self._call_method("mean")

        @derived_from(pd_Rolling)
        def median(self):
            return self._call_method("median")

        @derived_from(pd_Rolling)
        def min(self):
            return self._call_method("min")

        @derived_from(pd_Rolling)
        def max(self):
            return self._call_method("max")

        @derived_from(pd_Rolling)
        def std(self, ddof=1):
            return self._call_method("std", ddof=1)

        @derived_from(pd_Rolling)
        def var(self, ddof=1):
            return self._call_method("var", ddof=1)

        @derived_from(pd_Rolling)
        def skew(self):
            return self._call_method("skew")

        @derived_from(pd_Rolling)
        def kurt(self):
            return self._call_method("kurt")

        @derived_from(pd_Rolling)
        def quantile(self, quantile):
            return self._call_method("quantile", quantile)

        @derived_from(pd_Rolling)
        def apply(
            self,
            func,
            raw=False,
            engine="cython",
            engine_kwargs=None,
            args=None,
            kwargs=None,
        ):
            kwargs = kwargs or {}
            args = args or ()
            return self._call_method(
                "apply",
                func,
                raw=raw,
                engine=engine,
                engine_kwargs=engine_kwargs,
                args=args,
                kwargs=kwargs,
            )

        @derived_from(pd_Rolling)
        def aggregate(self, func, *args, **kwargs):
            return self._call_method("agg", func, *args, **kwargs)

        agg = aggregate

        def __repr__(self):
            def order(item):
                k, v = item
                _order = {
                    "window": 0,
                    "min_periods": 1,
                    "center": 2,
                    "win_type": 3,
                    "axis": 4,
                }
                return _order[k]

            rolling_kwargs = self._rolling_kwargs()
            rolling_kwargs["window"] = self.window
            rolling_kwargs["win_type"] = self._win_type
            return "Rolling [{}]".format(
                ",".join(
                    f"{k}={v}"
                    for k, v in sorted(rolling_kwargs.items(), key=order)
                    if v is not None
                )
            )

    class RollingGroupby(Rolling):
        def __init__(
            self,
            groupby,
            window=None,
            min_periods=None,
            center=False,
            win_type=None,
            axis=0,
        ):
            self._groupby_kwargs = groupby._groupby_kwargs
            self._groupby_slice = groupby._slice

            obj = groupby.obj
            if self._groupby_slice is not None:
                if isinstance(self._groupby_slice, str):
                    sliced_plus = [self._groupby_slice]
                else:
                    sliced_plus = list(self._groupby_slice)
                if isinstance(groupby.by, str):
                    sliced_plus.append(groupby.by)
                else:
                    sliced_plus.extend(groupby.by)
                obj = obj[sliced_plus]

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

        def _rolling_kwargs(self):
            kwargs = super()._rolling_kwargs()
            if kwargs.get("axis", None) in (0, "index"):
                # it's a default, no need to pass
                kwargs.pop("axis")
            return kwargs

        @staticmethod
        def pandas_rolling_method(
            df,
            rolling_kwargs,
            name,
            *args,
            groupby_kwargs=None,
            groupby_slice=None,
            **kwargs,
        ):
            groupby = df.groupby(**groupby_kwargs)
            if groupby_slice:
                groupby = groupby[groupby_slice]
            rolling = groupby.rolling(**rolling_kwargs)
            return getattr(rolling, name)(*args, **kwargs).sort_index(level=-1)

        def _call_method(self, method_name, *args, **kwargs):
            return super()._call_method(
                method_name,
                *args,
                groupby_kwargs=self._groupby_kwargs,
                groupby_slice=self._groupby_slice,
                **kwargs,
            )