from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd
import pytest

import xarray as xr
from xarray import DataArray, Dataset, set_options
from xarray.tests import (
    assert_allclose,
    assert_equal,
    assert_identical,
    has_dask,
    requires_dask,
    requires_dask_ge_2024_11_0,
    requires_numbagg,
)

pytestmark = [
    pytest.mark.filterwarnings("error:Mean of empty slice"),
    pytest.mark.filterwarnings("error:All-NaN (slice|axis) encountered"),
]


@pytest.mark.parametrize("func", ["mean", "sum"])
@pytest.mark.parametrize("min_periods", [1, 10])
def test_cumulative(d, func, min_periods) -> None:
    # One dim
    result = getattr(d.cumulative("z", min_periods=min_periods), func)()
    expected = getattr(d.rolling(z=d["z"].size, min_periods=min_periods), func)()
    assert_identical(result, expected)

    # Multiple dim
    result = getattr(d.cumulative(["z", "x"], min_periods=min_periods), func)()
    expected = getattr(
        d.rolling(z=d["z"].size, x=d["x"].size, min_periods=min_periods),
        func,
    )()
    assert_identical(result, expected)


def test_cumulative_vs_cum(d) -> None:
    result = d.cumulative("z").sum()
    expected = d.cumsum("z")
    # cumsum drops the coord of the dimension; cumulative doesn't
    expected = expected.assign_coords(z=result["z"])
    assert_identical(result, expected)


class TestDataArrayRolling:
    @pytest.mark.parametrize("da", (1, 2), indirect=True)
    @pytest.mark.parametrize("center", [True, False])
    @pytest.mark.parametrize("size", [1, 2, 3, 7])
    def test_rolling_iter(self, da: DataArray, center: bool, size: int) -> None:
        rolling_obj = da.rolling(time=size, center=center)
        rolling_obj_mean = rolling_obj.mean()

        assert len(rolling_obj.window_labels) == len(da["time"])
        assert_identical(rolling_obj.window_labels, da["time"])

        for i, (label, window_da) in enumerate(rolling_obj):
            assert label == da["time"].isel(time=i)

            actual = rolling_obj_mean.isel(time=i)
            expected = window_da.mean("time")

            np.testing.assert_allclose(actual.values, expected.values)

    @pytest.mark.parametrize("da", (1,), indirect=True)
    def test_rolling_repr(self, da) -> None:
        rolling_obj = da.rolling(time=7)
        assert repr(rolling_obj) == "DataArrayRolling [time->7]"
        rolling_obj = da.rolling(time=7, center=True)
        assert repr(rolling_obj) == "DataArrayRolling [time->7(center)]"
        rolling_obj = da.rolling(time=7, x=3, center=True)
        assert repr(rolling_obj) == "DataArrayRolling [time->7(center),x->3(center)]"

    @requires_dask
    def test_repeated_rolling_rechunks(self) -> None:
        # regression test for GH3277, GH2514
        dat = DataArray(np.random.rand(7653, 300), dims=("day", "item"))
        dat_chunk = dat.chunk({"item": 20})
        dat_chunk.rolling(day=10).mean().rolling(day=250).std()

    def test_rolling_doc(self, da) -> None:
        rolling_obj = da.rolling(time=7)

        # argument substitution worked
        assert "`mean`" in rolling_obj.mean.__doc__

    def test_rolling_properties(self, da) -> None:
        rolling_obj = da.rolling(time=4)

        assert rolling_obj.obj.get_axis_num("time") == 1

        # catching invalid args
        with pytest.raises(ValueError, match="window must be > 0"):
            da.rolling(time=-2)

        with pytest.raises(ValueError, match="min_periods must be greater than zero"):
            da.rolling(time=2, min_periods=0)

        with pytest.raises(
            KeyError,
            match=r"\('foo',\) not found in DataArray dimensions",
        ):
            da.rolling(foo=2)

    @requires_dask
    @pytest.mark.parametrize(
        "name", ("sum", "mean", "std", "min", "max", "median", "argmin", "argmax")
    )
    @pytest.mark.parametrize("center", (True, False, None))
    @pytest.mark.parametrize("min_periods", (1, None))
    @pytest.mark.parametrize("backend", ["numpy", "dask"], indirect=True)
    def test_rolling_wrapped_bottleneck(
        self, da, name, center, min_periods, compute_backend
    ) -> None:
        bn = pytest.importorskip("bottleneck", minversion="1.1")
        # Test all bottleneck functions
        rolling_obj = da.rolling(time=7, min_periods=min_periods)

        func_name = f"move_{name}"
        actual = getattr(rolling_obj, name)()
        window = 7
        expected = getattr(bn, func_name)(
            da.values, window=window, axis=1, min_count=min_periods
        )
        # index 0 is at the rightmost edge of the window
        # need to reverse index here
        # see GH #8541
        if func_name in ["move_argmin", "move_argmax"]:
            expected = window - 1 - expected

        # Using assert_allclose because we get tiny (1e-17) differences in numbagg.
        np.testing.assert_allclose(actual.values, expected)

        with pytest.warns(DeprecationWarning, match="Reductions are applied"):
            getattr(rolling_obj, name)(dim="time")

        # Test center
        rolling_obj = da.rolling(time=7, center=center)
        actual = getattr(rolling_obj, name)()["time"]
        # Using assert_allclose because we get tiny (1e-17) differences in numbagg.
        assert_allclose(actual, da["time"])

    @requires_dask
    @pytest.mark.parametrize("name", ("mean", "count"))
    @pytest.mark.parametrize("center", (True, False, None))
    @pytest.mark.parametrize("min_periods", (1, None))
    @pytest.mark.parametrize("window", (7, 8))
    @pytest.mark.parametrize("backend", ["dask"], indirect=True)
    def test_rolling_wrapped_dask(self, da, name, center, min_periods, window) -> None:
        # dask version
        rolling_obj = da.rolling(time=window, min_periods=min_periods, center=center)
        actual = getattr(rolling_obj, name)().load()
        if name != "count":
            with pytest.warns(DeprecationWarning, match="Reductions are applied"):
                getattr(rolling_obj, name)(dim="time")
        # numpy version
        rolling_obj = da.load().rolling(
            time=window, min_periods=min_periods, center=center
        )
        expected = getattr(rolling_obj, name)()

        # using all-close because rolling over ghost cells introduces some
        # precision errors
        assert_allclose(actual, expected)

        # with zero chunked array GH:2113
        rolling_obj = da.chunk().rolling(
            time=window, min_periods=min_periods, center=center
        )
        actual = getattr(rolling_obj, name)().load()
        assert_allclose(actual, expected)

    @pytest.mark.parametrize("center", (True, None))
    def test_rolling_wrapped_dask_nochunk(self, center) -> None:
        # GH:2113
        pytest.importorskip("dask.array")

        da_day_clim = xr.DataArray(
            np.arange(1, 367), coords=[np.arange(1, 367)], dims="dayofyear"
        )
        expected = da_day_clim.rolling(dayofyear=31, center=center).mean()
        actual = da_day_clim.chunk().rolling(dayofyear=31, center=center).mean()
        assert_allclose(actual, expected)

    @pytest.mark.parametrize("center", (True, False))
    @pytest.mark.parametrize("min_periods", (None, 1, 2, 3))
    @pytest.mark.parametrize("window", (1, 2, 3, 4))
    def test_rolling_pandas_compat(
        self, center, window, min_periods, compute_backend
    ) -> None:
        s = pd.Series(np.arange(10))
        da = DataArray.from_series(s)

        if min_periods is not None and window < min_periods:
            min_periods = window

        s_rolling = s.rolling(window, center=center, min_periods=min_periods).mean()
        da_rolling = da.rolling(
            index=window, center=center, min_periods=min_periods
        ).mean()
        da_rolling_np = da.rolling(
            index=window, center=center, min_periods=min_periods
        ).reduce(np.nanmean)

        np.testing.assert_allclose(np.asarray(s_rolling.values), da_rolling.values)
        np.testing.assert_allclose(s_rolling.index, da_rolling["index"])
        np.testing.assert_allclose(np.asarray(s_rolling.values), da_rolling_np.values)
        np.testing.assert_allclose(s_rolling.index, da_rolling_np["index"])

    @pytest.mark.parametrize("center", (True, False))
    @pytest.mark.parametrize("window", (1, 2, 3, 4))
    def test_rolling_construct(self, center: bool, window: int) -> None:
        s = pd.Series(np.arange(10))
        da = DataArray.from_series(s)

        s_rolling = s.rolling(window, center=center, min_periods=1).mean()
        da_rolling = da.rolling(index=window, center=center, min_periods=1)

        da_rolling_mean = da_rolling.construct("window").mean("window")
        np.testing.assert_allclose(np.asarray(s_rolling.values), da_rolling_mean.values)
        np.testing.assert_allclose(s_rolling.index, da_rolling_mean["index"])

        # with stride
        da_rolling_mean = da_rolling.construct("window", stride=2).mean("window")
        np.testing.assert_allclose(
            np.asarray(s_rolling.values[::2]), da_rolling_mean.values
        )
        np.testing.assert_allclose(s_rolling.index[::2], da_rolling_mean["index"])

        # with fill_value
        da_rolling_mean = da_rolling.construct("window", stride=2, fill_value=0.0).mean(
            "window"
        )
        assert da_rolling_mean.isnull().sum() == 0
        assert (da_rolling_mean == 0.0).sum() >= 0

    @pytest.mark.parametrize("da", (1, 2), indirect=True)
    @pytest.mark.parametrize("center", (True, False))
    @pytest.mark.parametrize("min_periods", (None, 1, 2, 3))
    @pytest.mark.parametrize("window", (1, 2, 3, 4))
    @pytest.mark.parametrize("name", ("sum", "mean", "std", "max"))
    def test_rolling_reduce(
        self, da, center, min_periods, window, name, compute_backend
    ) -> None:
        if min_periods is not None and window < min_periods:
            min_periods = window

        if da.isnull().sum() > 1 and window == 1:
            # this causes all nan slices
            window = 2

        rolling_obj = da.rolling(time=window, center=center, min_periods=min_periods)

        # add nan prefix to numpy methods to get similar # behavior as bottleneck
        actual = rolling_obj.reduce(getattr(np, f"nan{name}"))
        expected = getattr(rolling_obj, name)()
        assert_allclose(actual, expected)
        assert actual.sizes == expected.sizes

    @pytest.mark.parametrize("center", (True, False))
    @pytest.mark.parametrize("min_periods", (None, 1, 2, 3))
    @pytest.mark.parametrize("window", (1, 2, 3, 4))
    @pytest.mark.parametrize("name", ("sum", "max"))
    def test_rolling_reduce_nonnumeric(
        self, center, min_periods, window, name, compute_backend
    ) -> None:
        da = DataArray(
            [0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time"
        ).isnull()

        if min_periods is not None and window < min_periods:
            min_periods = window

        rolling_obj = da.rolling(time=window, center=center, min_periods=min_periods)

        # add nan prefix to numpy methods to get similar behavior as bottleneck
        actual = rolling_obj.reduce(getattr(np, f"nan{name}"))
        expected = getattr(rolling_obj, name)()
        assert_allclose(actual, expected)
        assert actual.sizes == expected.sizes

    def test_rolling_count_correct(self, compute_backend) -> None:
        da = DataArray([0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time")

        kwargs: list[dict[str, Any]] = [
            {"time": 11, "min_periods": 1},
            {"time": 11, "min_periods": None},
            {"time": 7, "min_periods": 2},
        ]
        expecteds = [
            DataArray([1, 1, 2, 3, 3, 4, 5, 6, 6, 7, 8], dims="time"),
            DataArray(
                [
                    np.nan,
                    np.nan,
                    np.nan,
                    np.nan,
                    np.nan,
                    np.nan,
                    np.nan,
                    np.nan,
                    np.nan,
                    np.nan,
                    np.nan,
                ],
                dims="time",
            ),
            DataArray([np.nan, np.nan, 2, 3, 3, 4, 5, 5, 5, 5, 5], dims="time"),
        ]

        for kwarg, expected in zip(kwargs, expecteds, strict=True):
            result = da.rolling(**kwarg).count()
            assert_equal(result, expected)

            result = da.to_dataset(name="var1").rolling(**kwarg).count()["var1"]
            assert_equal(result, expected)

    @pytest.mark.parametrize("da", (1,), indirect=True)
    @pytest.mark.parametrize("center", (True, False))
    @pytest.mark.parametrize("min_periods", (None, 1))
    @pytest.mark.parametrize("name", ("sum", "mean", "max"))
    def test_ndrolling_reduce(
        self, da, center, min_periods, name, compute_backend
    ) -> None:
        rolling_obj = da.rolling(time=3, x=2, center=center, min_periods=min_periods)

        actual = getattr(rolling_obj, name)()
        expected = getattr(
            getattr(
                da.rolling(time=3, center=center, min_periods=min_periods), name
            )().rolling(x=2, center=center, min_periods=min_periods),
            name,
        )()

        assert_allclose(actual, expected)
        assert actual.sizes == expected.sizes

        if name == "mean":
            # test our reimplementation of nanmean using np.nanmean
            expected = getattr(rolling_obj.construct({"time": "tw", "x": "xw"}), name)(
                ["tw", "xw"]
            )
            count = rolling_obj.count()
            if min_periods is None:
                min_periods = 1
            assert_allclose(actual, expected.where(count >= min_periods))

    @pytest.mark.parametrize("center", (True, False, (True, False)))
    @pytest.mark.parametrize("fill_value", (np.nan, 0.0))
    def test_ndrolling_construct(self, center, fill_value) -> None:
        da = DataArray(
            np.arange(5 * 6 * 7).reshape(5, 6, 7).astype(float),
            dims=["x", "y", "z"],
            coords={"x": ["a", "b", "c", "d", "e"], "y": np.arange(6)},
        )
        actual = da.rolling(x=3, z=2, center=center).construct(
            x="x1", z="z1", fill_value=fill_value
        )
        if not isinstance(center, tuple):
            center = (center, center)
        expected = (
            da.rolling(x=3, center=center[0])
            .construct(x="x1", fill_value=fill_value)
            .rolling(z=2, center=center[1])
            .construct(z="z1", fill_value=fill_value)
        )
        assert_allclose(actual, expected)

    @pytest.mark.parametrize(
        "funcname, argument",
        [
            ("reduce", (np.mean,)),
            ("mean", ()),
            ("construct", ("window_dim",)),
            ("count", ()),
        ],
    )
    def test_rolling_keep_attrs(self, funcname, argument) -> None:
        attrs_da = {"da_attr": "test"}

        data = np.linspace(10, 15, 100)
        coords = np.linspace(1, 10, 100)

        da = DataArray(
            data, dims=("coord"), coords={"coord": coords}, attrs=attrs_da, name="name"
        )

        # attrs are now kept per default
        func = getattr(da.rolling(dim={"coord": 5}), funcname)
        result = func(*argument)
        assert result.attrs == attrs_da
        assert result.name == "name"

        # discard attrs
        func = getattr(da.rolling(dim={"coord": 5}), funcname)
        result = func(*argument, keep_attrs=False)
        assert result.attrs == {}
        assert result.name == "name"

        # test discard attrs using global option
        func = getattr(da.rolling(dim={"coord": 5}), funcname)
        with set_options(keep_attrs=False):
            result = func(*argument)
        assert result.attrs == {}
        assert result.name == "name"

        # keyword takes precedence over global option
        func = getattr(da.rolling(dim={"coord": 5}), funcname)
        with set_options(keep_attrs=False):
            result = func(*argument, keep_attrs=True)
        assert result.attrs == attrs_da
        assert result.name == "name"

        func = getattr(da.rolling(dim={"coord": 5}), funcname)
        with set_options(keep_attrs=True):
            result = func(*argument, keep_attrs=False)
        assert result.attrs == {}
        assert result.name == "name"

    @requires_dask
    @pytest.mark.parametrize("dtype", ["int", "float32", "float64"])
    def test_rolling_dask_dtype(self, dtype) -> None:
        data = DataArray(
            np.array([1, 2, 3], dtype=dtype), dims="x", coords={"x": [1, 2, 3]}
        )
        unchunked_result = data.rolling(x=3, min_periods=1).mean()
        chunked_result = data.chunk({"x": 1}).rolling(x=3, min_periods=1).mean()
        assert chunked_result.dtype == unchunked_result.dtype

    def test_rolling_mean_bool(self) -> None:
        bool_raster = DataArray(
            data=[0, 1, 1, 0, 1, 0],
            dims=("x"),
        ).astype(bool)

        expected = DataArray(
            data=[np.nan, 2 / 3, 2 / 3, 2 / 3, 1 / 3, np.nan],
            dims=("x"),
        )

        result = bool_raster.rolling(x=3, center=True).mean()
        assert_allclose(result, expected)


@requires_numbagg
class TestDataArrayRollingExp:
    @pytest.mark.parametrize("dim", ["time", "x"])
    @pytest.mark.parametrize(
        "window_type, window",
        [["span", 5], ["alpha", 0.5], ["com", 0.5], ["halflife", 5]],
    )
    @pytest.mark.parametrize("backend", ["numpy"], indirect=True)
    @pytest.mark.parametrize("func", ["mean", "sum", "var", "std"])
    def test_rolling_exp_runs(self, da, dim, window_type, window, func) -> None:
        da = da.where(da > 0.2)

        rolling_exp = da.rolling_exp(window_type=window_type, **{dim: window})
        result = getattr(rolling_exp, func)()
        assert isinstance(result, DataArray)

    @pytest.mark.parametrize("dim", ["time", "x"])
    @pytest.mark.parametrize(
        "window_type, window",
        [["span", 5], ["alpha", 0.5], ["com", 0.5], ["halflife", 5]],
    )
    @pytest.mark.parametrize("backend", ["numpy"], indirect=True)
    def test_rolling_exp_mean_pandas(self, da, dim, window_type, window) -> None:
        da = da.isel(a=0).where(lambda x: x > 0.2)

        result = da.rolling_exp(window_type=window_type, **{dim: window}).mean()
        assert isinstance(result, DataArray)

        pandas_array = da.to_pandas()
        assert pandas_array.index.name == "time"
        if dim == "x":
            pandas_array = pandas_array.T
        expected = xr.DataArray(
            pandas_array.ewm(**{window_type: window}).mean()
        ).transpose(*da.dims)

        assert_allclose(expected.variable, result.variable)

    @pytest.mark.parametrize("backend", ["numpy"], indirect=True)
    @pytest.mark.parametrize("func", ["mean", "sum"])
    def test_rolling_exp_keep_attrs(self, da, func) -> None:
        attrs = {"attrs": "da"}
        da.attrs = attrs

        # Equivalent of `da.rolling_exp(time=10).mean`
        rolling_exp_func = getattr(da.rolling_exp(time=10), func)

        # attrs are kept per default
        result = rolling_exp_func()
        assert result.attrs == attrs

        # discard attrs
        result = rolling_exp_func(keep_attrs=False)
        assert result.attrs == {}

        # test discard attrs using global option
        with set_options(keep_attrs=False):
            result = rolling_exp_func()
        assert result.attrs == {}

        # keyword takes precedence over global option
        with set_options(keep_attrs=False):
            result = rolling_exp_func(keep_attrs=True)
        assert result.attrs == attrs

        with set_options(keep_attrs=True):
            result = rolling_exp_func(keep_attrs=False)
        assert result.attrs == {}

        with pytest.warns(
            UserWarning,
            match="Passing ``keep_attrs`` to ``rolling_exp`` has no effect.",
        ):
            da.rolling_exp(time=10, keep_attrs=True)


class TestDatasetRolling:
    @pytest.mark.parametrize(
        "funcname, argument",
        [
            ("reduce", (np.mean,)),
            ("mean", ()),
            ("construct", ("window_dim",)),
            ("count", ()),
        ],
    )
    def test_rolling_keep_attrs(self, funcname, argument) -> None:
        global_attrs = {"units": "test", "long_name": "testing"}
        da_attrs = {"da_attr": "test"}
        da_not_rolled_attrs = {"da_not_rolled_attr": "test"}

        data = np.linspace(10, 15, 100)
        coords = np.linspace(1, 10, 100)

        ds = Dataset(
            data_vars={"da": ("coord", data), "da_not_rolled": ("no_coord", data)},
            coords={"coord": coords},
            attrs=global_attrs,
        )
        ds.da.attrs = da_attrs
        ds.da_not_rolled.attrs = da_not_rolled_attrs

        # attrs are now kept per default
        func = getattr(ds.rolling(dim={"coord": 5}), funcname)
        result = func(*argument)
        assert result.attrs == global_attrs
        assert result.da.attrs == da_attrs
        assert result.da_not_rolled.attrs == da_not_rolled_attrs
        assert result.da.name == "da"
        assert result.da_not_rolled.name == "da_not_rolled"

        # discard attrs
        func = getattr(ds.rolling(dim={"coord": 5}), funcname)
        result = func(*argument, keep_attrs=False)
        assert result.attrs == {}
        assert result.da.attrs == {}
        assert result.da_not_rolled.attrs == {}
        assert result.da.name == "da"
        assert result.da_not_rolled.name == "da_not_rolled"

        # test discard attrs using global option
        func = getattr(ds.rolling(dim={"coord": 5}), funcname)
        with set_options(keep_attrs=False):
            result = func(*argument)

        assert result.attrs == {}
        assert result.da.attrs == {}
        assert result.da_not_rolled.attrs == {}
        assert result.da.name == "da"
        assert result.da_not_rolled.name == "da_not_rolled"

        # keyword takes precedence over global option
        func = getattr(ds.rolling(dim={"coord": 5}), funcname)
        with set_options(keep_attrs=False):
            result = func(*argument, keep_attrs=True)

        assert result.attrs == global_attrs
        assert result.da.attrs == da_attrs
        assert result.da_not_rolled.attrs == da_not_rolled_attrs
        assert result.da.name == "da"
        assert result.da_not_rolled.name == "da_not_rolled"

        func = getattr(ds.rolling(dim={"coord": 5}), funcname)
        with set_options(keep_attrs=True):
            result = func(*argument, keep_attrs=False)

        assert result.attrs == {}
        assert result.da.attrs == {}
        assert result.da_not_rolled.attrs == {}
        assert result.da.name == "da"
        assert result.da_not_rolled.name == "da_not_rolled"

    def test_rolling_properties(self, ds) -> None:
        # catching invalid args
        with pytest.raises(ValueError, match="window must be > 0"):
            ds.rolling(time=-2)
        with pytest.raises(ValueError, match="min_periods must be greater than zero"):
            ds.rolling(time=2, min_periods=0)
        with pytest.raises(KeyError, match="time2"):
            ds.rolling(time2=2)
        with pytest.raises(
            KeyError,
            match=r"\('foo',\) not found in Dataset dimensions",
        ):
            ds.rolling(foo=2)

    @requires_dask_ge_2024_11_0
    def test_rolling_construct_automatic_rechunk(self):
        import dask

        # Construct dataset with chunk size of (400, 400, 1) or 1.22 MiB
        da = DataArray(
            dims=["latitude", "longitude", "time"],
            data=dask.array.random.random((400, 400, 400), chunks=(-1, -1, 1)),
            name="foo",
        )

        for obj in [da, da.to_dataset()]:
            # Dataset now has chunks of size (400, 400, 100 100) or 11.92 GiB
            rechunked = obj.rolling(time=100, center=True).construct(
                "window",
                sliding_window_view_kwargs=dict(
                    automatic_rechunk=True, writeable=False
                ),
            )
            not_rechunked = obj.rolling(time=100, center=True).construct(
                "window",
                sliding_window_view_kwargs=dict(
                    automatic_rechunk=False, writeable=True
                ),
            )
            assert rechunked.chunksizes != not_rechunked.chunksizes

            roller = obj.isel(time=slice(30)).rolling(time=10, center=True)
            one = roller.reduce(
                np.sum, sliding_window_view_kwargs=dict(automatic_rechunk=True)
            )
            two = roller.reduce(
                np.sum, sliding_window_view_kwargs=dict(automatic_rechunk=False)
            )
            assert_identical(one, two)

    @pytest.mark.parametrize(
        "name", ("sum", "mean", "std", "var", "min", "max", "median")
    )
    @pytest.mark.parametrize("center", (True, False, None))
    @pytest.mark.parametrize("min_periods", (1, None))
    @pytest.mark.parametrize("key", ("z1", "z2"))
    @pytest.mark.parametrize("backend", ["numpy"], indirect=True)
    def test_rolling_wrapped_bottleneck(
        self, ds, name, center, min_periods, key, compute_backend
    ) -> None:
        bn = pytest.importorskip("bottleneck", minversion="1.1")

        # Test all bottleneck functions
        rolling_obj = ds.rolling(time=7, min_periods=min_periods)

        func_name = f"move_{name}"
        actual = getattr(rolling_obj, name)()
        if key == "z1":  # z1 does not depend on 'Time' axis. Stored as it is.
            expected = ds[key]
        elif key == "z2":
            expected = getattr(bn, func_name)(
                ds[key].values, window=7, axis=0, min_count=min_periods
            )
        else:
            raise ValueError
        np.testing.assert_allclose(actual[key].values, expected)

        # Test center
        rolling_obj = ds.rolling(time=7, center=center)
        actual = getattr(rolling_obj, name)()["time"]
        assert_allclose(actual, ds["time"])

    @pytest.mark.parametrize("center", (True, False))
    @pytest.mark.parametrize("min_periods", (None, 1, 2, 3))
    @pytest.mark.parametrize("window", (1, 2, 3, 4))
    def test_rolling_pandas_compat(self, center, window, min_periods) -> None:
        df = pd.DataFrame(
            {
                "x": np.random.randn(20),
                "y": np.random.randn(20),
                "time": np.linspace(0, 1, 20),
            }
        )
        ds = Dataset.from_dataframe(df)

        if min_periods is not None and window < min_periods:
            min_periods = window

        df_rolling = df.rolling(window, center=center, min_periods=min_periods).mean()
        ds_rolling = ds.rolling(
            index=window, center=center, min_periods=min_periods
        ).mean()

        np.testing.assert_allclose(
            np.asarray(df_rolling["x"].values), ds_rolling["x"].values
        )
        np.testing.assert_allclose(df_rolling.index, ds_rolling["index"])

    @pytest.mark.parametrize("center", (True, False))
    @pytest.mark.parametrize("window", (1, 2, 3, 4))
    def test_rolling_construct(self, center: bool, window: int) -> None:
        df = pd.DataFrame(
            {
                "x": np.random.randn(20),
                "y": np.random.randn(20),
                "time": np.linspace(0, 1, 20),
            }
        )

        ds = Dataset.from_dataframe(df)
        df_rolling = df.rolling(window, center=center, min_periods=1).mean()
        ds_rolling = ds.rolling(index=window, center=center)

        ds_rolling_mean = ds_rolling.construct("window").mean("window")
        np.testing.assert_allclose(
            np.asarray(df_rolling["x"].values), ds_rolling_mean["x"].values
        )
        np.testing.assert_allclose(df_rolling.index, ds_rolling_mean["index"])

        # with fill_value
        ds_rolling_mean = ds_rolling.construct("window", stride=2, fill_value=0.0).mean(
            "window"
        )
        assert (ds_rolling_mean.isnull().sum() == 0).to_dataarray(dim="vars").all()
        assert (ds_rolling_mean["x"] == 0.0).sum() >= 0

    @pytest.mark.parametrize("center", (True, False))
    @pytest.mark.parametrize("window", (1, 2, 3, 4))
    def test_rolling_construct_stride(self, center: bool, window: int) -> None:
        df = pd.DataFrame(
            {
                "x": np.random.randn(20),
                "y": np.random.randn(20),
                "time": np.linspace(0, 1, 20),
            }
        )
        ds = Dataset.from_dataframe(df)
        df_rolling_mean = df.rolling(window, center=center, min_periods=1).mean()

        # With an index (dimension coordinate)
        ds_rolling = ds.rolling(index=window, center=center)
        ds_rolling_mean = ds_rolling.construct("w", stride=2).mean("w")
        np.testing.assert_allclose(
            np.asarray(df_rolling_mean["x"][::2].values), ds_rolling_mean["x"].values
        )
        np.testing.assert_allclose(df_rolling_mean.index[::2], ds_rolling_mean["index"])

        # Without index (https://github.com/pydata/xarray/issues/7021)
        ds2 = ds.drop_vars("index")
        ds2_rolling = ds2.rolling(index=window, center=center)
        ds2_rolling_mean = ds2_rolling.construct("w", stride=2).mean("w")
        np.testing.assert_allclose(
            np.asarray(df_rolling_mean["x"][::2].values), ds2_rolling_mean["x"].values
        )

        # Mixed coordinates, indexes and 2D coordinates
        ds3 = xr.Dataset(
            {"x": ("t", range(20)), "x2": ("y", range(5))},
            {
                "t": range(20),
                "y": ("y", range(5)),
                "t2": ("t", range(20)),
                "y2": ("y", range(5)),
                "yt": (["t", "y"], np.ones((20, 5))),
            },
        )
        ds3_rolling = ds3.rolling(t=window, center=center)
        ds3_rolling_mean = ds3_rolling.construct("w", stride=2).mean("w")
        for coord in ds3.coords:
            assert coord in ds3_rolling_mean.coords

    @pytest.mark.slow
    @pytest.mark.parametrize("ds", (1, 2), indirect=True)
    @pytest.mark.parametrize("center", (True, False))
    @pytest.mark.parametrize("min_periods", (None, 1, 2, 3))
    @pytest.mark.parametrize("window", (1, 2, 3, 4))
    @pytest.mark.parametrize(
        "name", ("sum", "mean", "std", "var", "min", "max", "median")
    )
    def test_rolling_reduce(self, ds, center, min_periods, window, name) -> None:
        if min_periods is not None and window < min_periods:
            min_periods = window

        if name == "std" and window == 1:
            pytest.skip("std with window == 1 is unstable in bottleneck")

        rolling_obj = ds.rolling(time=window, center=center, min_periods=min_periods)

        # add nan prefix to numpy methods to get similar behavior as bottleneck
        actual = rolling_obj.reduce(getattr(np, f"nan{name}"))
        expected = getattr(rolling_obj, name)()
        assert_allclose(actual, expected)
        assert ds.sizes == actual.sizes
        # make sure the order of data_var are not changed.
        assert list(ds.data_vars.keys()) == list(actual.data_vars.keys())

        # Make sure the dimension order is restored
        for key, src_var in ds.data_vars.items():
            assert src_var.dims == actual[key].dims

    @pytest.mark.parametrize("ds", (2,), indirect=True)
    @pytest.mark.parametrize("center", (True, False))
    @pytest.mark.parametrize("min_periods", (None, 1))
    @pytest.mark.parametrize("name", ("sum", "max"))
    @pytest.mark.parametrize("dask", (True, False))
    def test_ndrolling_reduce(self, ds, center, min_periods, name, dask) -> None:
        if dask and has_dask:
            ds = ds.chunk({"x": 4})

        rolling_obj = ds.rolling(time=4, x=3, center=center, min_periods=min_periods)

        actual = getattr(rolling_obj, name)()
        expected = getattr(
            getattr(
                ds.rolling(time=4, center=center, min_periods=min_periods), name
            )().rolling(x=3, center=center, min_periods=min_periods),
            name,
        )()
        assert_allclose(actual, expected)
        assert actual.sizes == expected.sizes

        # Do it in the opposite order
        expected = getattr(
            getattr(
                ds.rolling(x=3, center=center, min_periods=min_periods), name
            )().rolling(time=4, center=center, min_periods=min_periods),
            name,
        )()

        assert_allclose(actual, expected)
        assert actual.sizes == expected.sizes

    @pytest.mark.parametrize("center", (True, False, (True, False)))
    @pytest.mark.parametrize("fill_value", (np.nan, 0.0))
    @pytest.mark.parametrize("dask", (True, False))
    def test_ndrolling_construct(self, center, fill_value, dask) -> None:
        da = DataArray(
            np.arange(5 * 6 * 7).reshape(5, 6, 7).astype(float),
            dims=["x", "y", "z"],
            coords={"x": ["a", "b", "c", "d", "e"], "y": np.arange(6)},
        )
        ds = xr.Dataset({"da": da})
        if dask and has_dask:
            ds = ds.chunk({"x": 4})

        actual = ds.rolling(x=3, z=2, center=center).construct(
            x="x1", z="z1", fill_value=fill_value
        )
        if not isinstance(center, tuple):
            center = (center, center)
        expected = (
            ds.rolling(x=3, center=center[0])
            .construct(x="x1", fill_value=fill_value)
            .rolling(z=2, center=center[1])
            .construct(z="z1", fill_value=fill_value)
        )
        assert_allclose(actual, expected)

    @requires_dask
    @pytest.mark.filterwarnings("error")
    @pytest.mark.parametrize("ds", (2,), indirect=True)
    @pytest.mark.parametrize("name", ("mean", "max"))
    def test_raise_no_warning_dask_rolling_assert_close(self, ds, name) -> None:
        """
        This is a puzzle — I can't easily find the source of the warning. It
        requires `assert_allclose` to be run, for the `ds` param to be 2, and is
        different for `mean` and `max`. `sum` raises no warning.
        """

        ds = ds.chunk({"x": 4})

        rolling_obj = ds.rolling(time=4, x=3)

        actual = getattr(rolling_obj, name)()
        expected = getattr(getattr(ds.rolling(time=4), name)().rolling(x=3), name)()
        assert_allclose(actual, expected)


@requires_numbagg
class TestDatasetRollingExp:
    @pytest.mark.parametrize(
        "backend", ["numpy", pytest.param("dask", marks=requires_dask)], indirect=True
    )
    def test_rolling_exp(self, ds) -> None:
        result = ds.rolling_exp(time=10, window_type="span").mean()
        assert isinstance(result, Dataset)

    @pytest.mark.parametrize("backend", ["numpy"], indirect=True)
    def test_rolling_exp_keep_attrs(self, ds) -> None:
        attrs_global = {"attrs": "global"}
        attrs_z1 = {"attr": "z1"}

        ds.attrs = attrs_global
        ds.z1.attrs = attrs_z1

        # attrs are kept per default
        result = ds.rolling_exp(time=10).mean()
        assert result.attrs == attrs_global
        assert result.z1.attrs == attrs_z1

        # discard attrs
        result = ds.rolling_exp(time=10).mean(keep_attrs=False)
        assert result.attrs == {}
        # TODO: from #8114 — this arguably should be empty, but `apply_ufunc` doesn't do
        # that at the moment. We should change in `apply_func` rather than
        # special-case it here.
        #
        # assert result.z1.attrs == {}

        # test discard attrs using global option
        with set_options(keep_attrs=False):
            result = ds.rolling_exp(time=10).mean()
        assert result.attrs == {}
        # See above
        # assert result.z1.attrs == {}

        # keyword takes precedence over global option
        with set_options(keep_attrs=False):
            result = ds.rolling_exp(time=10).mean(keep_attrs=True)
        assert result.attrs == attrs_global
        assert result.z1.attrs == attrs_z1

        with set_options(keep_attrs=True):
            result = ds.rolling_exp(time=10).mean(keep_attrs=False)
        assert result.attrs == {}
        # See above
        # assert result.z1.attrs == {}

        with pytest.warns(
            UserWarning,
            match="Passing ``keep_attrs`` to ``rolling_exp`` has no effect.",
        ):
            ds.rolling_exp(time=10, keep_attrs=True)