# ruff: noqa: E402
import pytest


pytest.importorskip("xarray")

from pytensor.xtensor.type import xtensor
from tests.xtensor.util import xr_arange_like, xr_assert_allclose, xr_function


@pytest.mark.parametrize(
    "dim", [..., None, "a", ("c", "a")], ids=["Ellipsis", "None", "a", "(a, c)"]
)
@pytest.mark.parametrize(
    "method", ["sum", "prod", "all", "any", "max", "min", "cumsum", "cumprod"][2:]
)
def test_reduction(method, dim):
    x = xtensor("x", dims=("a", "b", "c"), shape=(3, 5, 7))
    out = getattr(x, method)(dim=dim)

    fn = xr_function([x], out)
    x_test = xr_arange_like(x)

    xr_assert_allclose(
        fn(x_test),
        getattr(x_test, method)(dim=dim),
    )
