import itertools
import math
import re
import tracemalloc
from copy import copy

import numpy as np
import pytest

import pytensor
import pytensor.scalar as ps
import pytensor.tensor as pt
import tests.unittest_tools as utt
from pytensor import In, Out
from pytensor.compile.function import function
from pytensor.compile.mode import Mode
from pytensor.configdefaults import config
from pytensor.graph.basic import Apply, Variable
from pytensor.graph.fg import FunctionGraph
from pytensor.graph.replace import vectorize_node
from pytensor.link.basic import PerformLinker
from pytensor.link.c.basic import CLinker, OpWiseCLinker
from pytensor.npy_2_compat import numpy_maxdims
from pytensor.tensor import as_tensor_variable
from pytensor.tensor.basic import get_scalar_constant_value, second
from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
from pytensor.tensor.math import Any, Sum, exp
from pytensor.tensor.math import all as pt_all
from pytensor.tensor.math import any as pt_any
from pytensor.tensor.math import sum as pt_sum
from pytensor.tensor.type import (
    TensorType,
    bmatrix,
    bscalar,
    discrete_dtypes,
    lscalar,
    matrix,
    scalar,
    tensor,
    tensor3,
    vector,
    vectors,
)
from tests import unittest_tools
from tests.link.test_link import make_function


def reduce_bitwise_and(x, axis=-1, dtype="int8"):
    """Helper function for TestCAReduce"""
    if dtype == "uint8":
        # in numpy version >= 2.0, out of bounds uint8 values are not converted
        identity = np.array((255,), dtype=dtype)[0]
    else:
        identity = np.array((-1,), dtype=dtype)[0]

    shape_without_axis = tuple(s for i, s in enumerate(x.shape) if i != axis)
    if 0 in shape_without_axis:
        return np.empty(shape=shape_without_axis, dtype=x.dtype)

    def custom_reduce(a):
        out = identity
        for i in range(a.size):
            out = np.bitwise_and(a[i], out)
        return out

    return np.apply_along_axis(custom_reduce, axis, x)


def dimshuffle_benchmark(mode, c_contiguous, benchmark):
    x = tensor3("x")
    if c_contiguous:
        x_val = np.random.random((2, 3, 4)).astype(config.floatX)
    else:
        x_val = np.random.random((200, 300, 400)).transpose(1, 2, 0)
    ys = [x.transpose(t) for t in itertools.permutations((0, 1, 2))]
    ys += [
        x[None],
        x[:, None],
        x[:, :, None],
        x[:, :, :, None],
    ]
    # Borrow to avoid deepcopy overhead
    fn = pytensor.function(
        [In(x, borrow=True)],
        [Out(y, borrow=True) for y in ys],
        mode=mode,
    )
    fn.trust_input = True
    fn(x_val)  # JIT compile for JIT backends
    benchmark(fn, x_val)


class TestDimShuffle(unittest_tools.InferShapeTester):
    op = DimShuffle
    type = TensorType
    dtype = pytensor.config.floatX

    def with_linker(self, linker):
        for xsh, shuffle, zsh in [
            ((2, 3), (1, "x", 0), (3, 1, 2)),
            ((1, 2, 3), (1, 2), (2, 3)),
            ((1, 2, 1, 3), (1, 3), (2, 3)),
            ((2, 3, 4), (2, 1, 0), (4, 3, 2)),
            ((2, 3, 4), ("x", 2, 1, 0, "x"), (1, 4, 3, 2, 1)),
            ((1, 4, 3, 2, 1), (3, 2, 1), (2, 3, 4)),
            ((1, 1, 4), (1, 2), (1, 4)),
            ((1, 1, 1), (), ()),
            ((1,), ("x", "x"), (1, 1)),
        ]:
            i_shape = [entry if entry == 1 else None for entry in xsh]
            x = self.type(self.dtype, shape=i_shape)("x")
            e = self.op(input_ndim=len(i_shape), new_order=shuffle)(x)
            f = pytensor.function([x], e, mode=Mode(linker=linker))
            assert f(np.ones(xsh, dtype=self.dtype)).shape == zsh
            # test that DimShuffle.infer_shape work correctly
            x = self.type(self.dtype, shape=i_shape)("x")
            e = self.op(input_ndim=len(i_shape), new_order=shuffle)(x)
            f = pytensor.function(
                [x], e.shape, mode=Mode(linker=linker), on_unused_input="ignore"
            )
            assert all(f(np.ones(xsh, dtype=self.dtype))) == all(zsh)

        # Test when we drop a axis that is not broadcastable
        x = self.type(self.dtype, shape=(2, 1, None))("x")
        with pytest.raises(TypeError):
            self.op(input_ndim=3, new_order=shuffle)(x)

        # Test when we drop a axis that don't have shape 1
        x = self.type(self.dtype, shape=(None, 1, None))("x")
        e = self.op(input_ndim=3, new_order=(1, 2))(x)
        f = pytensor.function([x], e, mode=Mode(linker=linker))
        with pytest.raises(ValueError):
            f(np.ones((2, 1, 4), dtype=self.dtype))

        # Test that we can't take a dimensions multiple time
        xsh, shuffle, zsh = ((1, 1, 4), (0, 1, 2, 0), (1, 4))
        x = self.type(self.dtype, shape=(None, 1, None))("x")
        with pytest.raises(ValueError):
            DimShuffle(input_ndim=3, new_order=shuffle)

    def test_perform(self):
        self.with_linker(PerformLinker())

    def test_c_or_py(self):
        self.with_linker(OpWiseCLinker())

    def test_infer_shape(self):
        for xsh, shuffle in [
            ((2, 3), (1, "x", 0)),
            ((1, 2, 3), (1, 2)),
            ((1, 2, 1, 3), (1, 3)),
            ((2, 3, 4), (2, 1, 0)),
            ((2, 3, 4), ("x", 2, 1, 0, "x")),
            ((1, 4, 3, 2, 1), (3, 2, 1)),
            ((1, 1, 4), (1, 2)),
            ((1, 1, 1), ()),
            ((1,), ("x", "x")),
        ]:
            i_shape = [entry if entry == 1 else None for entry in xsh]
            adtens = self.type(self.dtype, shape=i_shape)("x")
            adtens_val = np.ones(xsh, dtype=self.dtype)
            self._compile_and_check(
                [adtens],
                [self.op(input_ndim=len(xsh), new_order=shuffle)(adtens)],
                [adtens_val],
                self.op,
                warn=False,
            )

    def test_too_big_rank(self):
        x = self.type(self.dtype, shape=())()
        y = x.dimshuffle(("x",) * (numpy_maxdims + 1))

        with pytest.raises(ValueError):
            y.eval({x: 0})

    def test_c_views(self):
        x_pt = vector()
        thunk, inputs, outputs = (
            CLinker().accept(FunctionGraph([x_pt], [x_pt[None]])).make_thunk()
        )

        # This is a little hackish, but we're hoping that--by running this more than
        # a few times--we're more likely to run into random memory that isn't the same
        # as the broadcasted value; that way, we'll be able to tell that we're getting
        # junk data from a poorly constructed array view.
        x_val = np.broadcast_to(2039, (5000,))
        expected_x_val = x_val[None]
        for i in range(1):
            inputs[0].storage[0] = x_val
            thunk()
            # Make sure it's a view of the original data
            assert np.shares_memory(x_val, outputs[0].storage[0])
            # Confirm the right strides
            assert outputs[0].storage[0].strides == expected_x_val.strides
            # Confirm the broadcasted value in the output
            assert np.array_equiv(outputs[0].storage[0], 2039)

    @pytest.mark.parametrize("inplace", [True, False])
    def test_memory_leak(self, inplace):
        import gc

        n = 100_000

        x = pytensor.shared(np.ones(n, dtype=np.float64))

        y = x.dimshuffle([0, "x"])
        y.owner.op.inplace = inplace

        f = pytensor.function([], y, mode=Mode(optimizer=None))

        assert len(f.maker.fgraph.apply_nodes) == 2
        assert isinstance(f.maker.fgraph.toposort()[0].op, DimShuffle)

        assert f.maker.fgraph.toposort()[0].op.inplace is inplace

        tracemalloc.start()

        blocks_last = None
        block_diffs = []
        for i in range(50):
            x.set_value(np.ones(n))
            _ = f()
            _ = gc.collect()
            blocks_i, _ = tracemalloc.get_traced_memory()
            if blocks_last is not None:
                blocks_diff = (blocks_i - blocks_last) // 10**3
                block_diffs.append(blocks_diff)
            blocks_last = blocks_i

        tracemalloc.stop()
        assert np.allclose(np.mean(block_diffs), 0)

    def test_static_shape(self):
        x = tensor(dtype=np.float64, shape=(1, 2), name="x")
        y = x.dimshuffle([0, 1, "x"])
        assert y.type.shape == (1, 2, 1)

    def test_valid_input_ndim(self):
        assert DimShuffle(input_ndim=2, new_order=(1, 0)).input_ndim == 2

        with pytest.raises(TypeError, match="input_ndim must be an integer"):
            DimShuffle(input_ndim=(True, False), new_order=(1, 0))

    @pytest.mark.parametrize("c_contiguous", [True, False])
    def test_benchmark(self, c_contiguous, benchmark):
        dimshuffle_benchmark("FAST_RUN", c_contiguous, benchmark)


class TestBroadcast:
    # this is to allow other types to reuse this class to test their ops
    type = TensorType
    op = Elemwise

    ctype = TensorType
    cop = Elemwise

    openmp_minsize = 2 * config.openmp_elemwise_minsize
    openmp_minsize_sqrt = int(math.ceil(math.sqrt(openmp_minsize)))

    # The order is important if you change them.
    linkers = [PerformLinker, CLinker]

    def rand_val(self, shp):
        return np.asarray(np.random.random(shp), dtype=pytensor.config.floatX)

    def rand_cval(self, shp):
        return np.asarray(np.random.random(shp), dtype=pytensor.config.floatX)

    def with_linker(self, linker, op, type, rand_val):
        for xsh, ysh in [
            ((3, 5), (3, 5)),
            ((3, 5), (1, 5)),
            ((3, 5), (3, 1)),
            ((1, 5), (5, 1)),
            ((1, 1), (1, 1)),
            ((self.openmp_minsize,), (self.openmp_minsize,)),
            (
                (self.openmp_minsize_sqrt, self.openmp_minsize_sqrt),
                (self.openmp_minsize_sqrt, self.openmp_minsize_sqrt),
            ),
            ((2, 3, 4, 5), (2, 3, 4, 5)),
            ((2, 3, 4, 5), (1, 3, 1, 5)),
            ((2, 3, 4, 5), (1, 1, 1, 1)),
            ((), ()),
        ]:
            x_type = type(
                pytensor.config.floatX,
                shape=tuple(s if s == 1 else None for s in xsh),
            )
            y_type = type(
                pytensor.config.floatX,
                shape=tuple(s if s == 1 else None for s in ysh),
            )

            x = x_type("x")
            y = y_type("y")
            e = op(ps.add)(x, y)
            f = make_function(copy(linker).accept(FunctionGraph([x, y], [e])))
            xv = rand_val(xsh)
            yv = rand_val(ysh)
            zv = xv + yv

            unittest_tools.assert_allclose(f(xv, yv), zv)

            # test Elemwise.infer_shape
            # the Shape op don't implement c_code!
            if isinstance(linker, PerformLinker):
                x = x_type("x")
                y = y_type("y")
                e = op(ps.add)(x, y)
                f = make_function(copy(linker).accept(FunctionGraph([x, y], [e.shape])))
                assert tuple(f(xv, yv)) == tuple(zv.shape)

    def with_linker_inplace(self, linker, op, type, rand_val):
        for xsh, ysh in [
            ((5, 5), (5, 5)),
            ((5, 5), (1, 5)),
            ((5, 5), (5, 1)),
            ((1, 1), (1, 1)),
            ((2, 3, 4, 5), (2, 3, 4, 5)),
            ((2, 3, 4, 5), (1, 3, 1, 5)),
            ((2, 3, 4, 5), (1, 1, 1, 1)),
            ((), ()),
        ]:
            x_type = type(
                pytensor.config.floatX,
                shape=tuple(s if s == 1 else None for s in xsh),
            )
            y_type = type(
                pytensor.config.floatX,
                shape=tuple(s if s == 1 else None for s in ysh),
            )

            x = x_type("x")
            y = y_type("y")
            e = op(ps.Add(ps.transfer_type(0)), {0: 0})(x, y)
            f = make_function(copy(linker).accept(FunctionGraph([x, y], [e])))
            xv = rand_val(xsh)
            yv = rand_val(ysh)
            zv = xv + yv

            f(xv, yv)

            assert (xv == zv).all()
            # test Elemwise.infer_shape
            # the Shape op don't implement c_code!
            if isinstance(linker, PerformLinker):
                x = x_type("x")
                y = y_type("y")
                e = op(ps.Add(ps.transfer_type(0)), {0: 0})(x, y)
                f = make_function(copy(linker).accept(FunctionGraph([x, y], [e.shape])))
                xv = rand_val(xsh)
                yv = rand_val(ysh)
                zv = xv + yv
                assert xv.shape == zv.shape
                assert tuple(f(xv, yv)) == zv.shape

    def test_perform(self):
        self.with_linker(PerformLinker(), self.op, self.type, self.rand_val)

    @pytest.mark.skipif(
        not pytensor.config.cxx,
        reason="G++ not available, so we need to skip this test.",
    )
    def test_c(self):
        self.with_linker(CLinker(), self.cop, self.ctype, self.rand_cval)

    def test_perform_inplace(self):
        self.with_linker_inplace(PerformLinker(), self.op, self.type, self.rand_val)

    @pytest.mark.skipif(
        not pytensor.config.cxx,
        reason="G++ not available, so we need to skip this test.",
    )
    def test_c_inplace(self):
        self.with_linker_inplace(CLinker(), self.cop, self.ctype, self.rand_cval)

    @pytest.mark.skipif(
        not pytensor.config.cxx,
        reason="G++ not available, so we need to skip this test.",
    )
    def test_fill(self):
        for linker, op, t, rval in zip(
            self.linkers,
            [self.op, self.cop],
            [self.type, self.ctype],
            [self.rand_val, self.rand_cval],
            strict=True,
        ):
            x = t(pytensor.config.floatX, shape=(None, None))("x")
            y = t(pytensor.config.floatX, shape=(1, 1))("y")
            e = op(ps.Second(ps.transfer_type(0)), {0: 0})(x, y)
            f = make_function(linker().accept(FunctionGraph([x, y], [e])))
            xv = rval((5, 5))
            yv = rval((1, 1))
            f(xv, yv)
            assert (xv == yv).all()

    def test_fill_var(self):
        x = matrix()
        x.fill(3)

    def test_fill_grad(self):
        x = TensorType(config.floatX, shape=(None, 1, None))("x")
        y = TensorType(config.floatX, shape=(None, 1, None))("y")
        e = second(x, y)
        # TODO FIXME: Make this a real test and assert something here!
        pytensor.grad(e.sum(), y)

    @pytest.mark.skipif(
        not pytensor.config.cxx,
        reason="G++ not available, so we need to skip this test.",
    )
    def test_weird_strides(self):
        for linker, op, t, rval in zip(
            self.linkers,
            [self.op, self.cop],
            [self.type, self.ctype],
            [self.rand_val, self.rand_cval],
            strict=True,
        ):
            x = t(pytensor.config.floatX, shape=(None,) * 5)("x")
            y = t(pytensor.config.floatX, shape=(None,) * 5)("y")
            e = op(ps.add)(x, y)
            f = make_function(linker().accept(FunctionGraph([x, y], [e])))
            xv = rval((2, 2, 2, 2, 2))
            yv = rval((2, 2, 2, 2, 2)).transpose(4, 0, 3, 1, 2)
            zv = xv + yv
            assert (f(xv, yv) == zv).all()

    @pytest.mark.skipif(
        not pytensor.config.cxx,
        reason="G++ not available, so we need to skip this test.",
    )
    def test_same_inputs(self):
        for linker, op, t, rval in zip(
            self.linkers,
            [self.op, self.cop],
            [self.type, self.ctype],
            [self.rand_val, self.rand_cval],
            strict=True,
        ):
            x = t(pytensor.config.floatX, shape=(None,) * 2)("x")
            e = op(ps.add)(x, x)
            f = make_function(linker().accept(FunctionGraph([x], [e])))
            xv = rval((2, 2))
            zv = xv + xv
            assert (f(xv) == zv).all()


class TestCAReduce(unittest_tools.InferShapeTester):
    op = CAReduce
    cases = [
        ((5, 6), None),
        ((5, 6), (0, 1)),
        ((5, 6), (0,)),
        ((5, 6), (1,)),
        ((5, 6), (-1,)),
        ((5, 6), (-2,)),
        ((5, 6), ()),
        ((2, 3, 4, 5), (0, 1, 3)),
        ((2, 3, 4, 5), (-2, -3)),
        ((5, 0), None),
        ((5, 0), (0,)),
        ((5, 0), (1,)),
        ((5, 0), ()),
        ((), None),
        ((), ()),
    ]
    type = TensorType

    def with_mode(
        self,
        mode,
        scalar_op=ps.add,
        dtype="floatX",
        pre_scalar_op=None,
        test_nan=False,
        tensor_op=None,
    ):
        for xsh, tosum in self.cases:
            if dtype == "floatX":
                dtype = pytensor.config.floatX
            x = self.type(
                dtype, shape=tuple(entry if entry == 1 else None for entry in xsh)
            )("x")
            d = {}
            if pre_scalar_op is not None:
                d = {"pre_scalar_op": pre_scalar_op}
            if tensor_op is None:
                e = as_tensor_variable(self.op(scalar_op, axis=tosum, **d)(x))
            else:
                e = as_tensor_variable(tensor_op(x, axis=tosum, **d))

            if tosum is None:
                tosum = list(range(len(xsh)))

            f = pytensor.function([x], e, mode=mode, on_unused_input="ignore")
            xv = np.asarray(np.random.random(xsh))

            if dtype not in discrete_dtypes:
                xv = np.asarray(xv, dtype=dtype)
            else:
                xv = np.asarray(xv < 0.5, dtype=dtype)

            if test_nan and xv.size > 0:
                if len(xsh) > 0:
                    xv = xv.flatten()
                    xv[0] = np.nan
                    xv = xv.reshape(*xsh)
                else:
                    xv = np.asarray(np.nan, dtype=dtype)
            zv = xv
            if pre_scalar_op is not None:
                zv = Elemwise(scalar_op=pre_scalar_op)(x).eval({x: xv})

            if len(tosum) > 1 and any(a < 0 for a in tosum):
                # In that case, we need to use the good order of axis
                # in the reduction.
                axis2 = []
                for a in tosum:
                    if a < 0:
                        axis2.append(a + len(xsh))
                    else:
                        axis2.append(a)
                assert len(axis2) == len(tosum)
                tosum = tuple(axis2)
            if tensor_op == pt_all:
                for axis in sorted(tosum, reverse=True):
                    zv = np.all(zv, axis)
                if len(tosum) == 0:
                    zv = zv != 0
            elif tensor_op == pt_any:
                for axis in sorted(tosum, reverse=True):
                    zv = np.any(zv, axis)
                if len(tosum) == 0:
                    zv = zv != 0
            elif scalar_op == ps.add:
                for axis in sorted(tosum, reverse=True):
                    zv = np.add.reduce(zv, axis)
                if dtype == "bool":
                    # np.add of a bool upcast, while CAReduce don't
                    zv = zv.astype(dtype)
            elif scalar_op == ps.mul:
                for axis in sorted(tosum, reverse=True):
                    zv = np.multiply.reduce(zv, axis)
            elif scalar_op == ps.scalar_maximum:
                # There is no identity value for the maximum function
                # So we can't support shape of dimensions 0.
                if np.prod(zv.shape) == 0:
                    continue
                for axis in sorted(tosum, reverse=True):
                    zv = np.maximum.reduce(zv, axis)
            elif scalar_op == ps.scalar_minimum:
                # There is no identity value for the minimum function
                # So we can't support shape of dimensions 0.
                if np.prod(zv.shape) == 0:
                    continue
                for axis in sorted(tosum, reverse=True):
                    zv = np.minimum.reduce(zv, axis)
            elif scalar_op == ps.or_:
                for axis in sorted(tosum, reverse=True):
                    zv = np.bitwise_or.reduce(zv, axis)
            elif scalar_op == ps.and_:
                for axis in sorted(tosum, reverse=True):
                    zv = reduce_bitwise_and(zv, axis, dtype=dtype)
            elif scalar_op == ps.xor:
                # There is no identity value for the xor function
                # So we can't support shape of dimensions 0.
                if np.prod(zv.shape) == 0:
                    continue
                for axis in sorted(tosum, reverse=True):
                    zv = np.bitwise_xor.reduce(zv, axis)
            else:
                raise NotImplementedError(
                    f"Test for CAReduce with scalar_op {scalar_op} not implemented"
                )

            if test_nan:
                assert self.type.values_eq(f(xv), zv), (f(xv), zv)
            else:
                f_xv = f(xv)
                assert f_xv.shape == zv.shape, (f_xv, zv)
                utt.assert_allclose(zv, f_xv)

            x = self.type(
                dtype, shape=tuple(entry if entry == 1 else None for entry in xsh)
            )("x")
            if tensor_op is None:
                e = self.op(scalar_op, axis=tosum)(x)
            else:
                e = tensor_op(x, axis=tosum)
            if tosum is None:
                tosum = list(range(len(xsh)))
            f = pytensor.function([x], e.shape, mode=mode, on_unused_input="ignore")
            if not (
                scalar_op in [ps.scalar_maximum, ps.scalar_minimum]
                and (xsh == () or np.prod(xsh) == 0)
            ):
                assert all(f(xv) == zv.shape)

    def test_perform_noopt(self):
        self.with_mode(Mode(linker="py", optimizer=None), ps.add, dtype="floatX")

    def test_perform(self):
        for dtype in ["bool", "floatX", "complex64", "complex128", "int8", "uint8"]:
            self.with_mode(Mode(linker="py"), ps.add, dtype=dtype)
            self.with_mode(Mode(linker="py"), ps.mul, dtype=dtype)
            self.with_mode(Mode(linker="py"), ps.scalar_maximum, dtype=dtype)
            self.with_mode(Mode(linker="py"), ps.scalar_minimum, dtype=dtype)
            self.with_mode(Mode(linker="py"), ps.and_, dtype=dtype, tensor_op=pt_all)
            self.with_mode(Mode(linker="py"), ps.or_, dtype=dtype, tensor_op=pt_any)
        for dtype in ["int8", "uint8"]:
            self.with_mode(Mode(linker="py"), ps.or_, dtype=dtype)
            self.with_mode(Mode(linker="py"), ps.and_, dtype=dtype)
            self.with_mode(Mode(linker="py"), ps.xor, dtype=dtype)

    def test_perform_nan(self):
        for dtype in ["floatX", "complex64", "complex128"]:
            self.with_mode(Mode(linker="py"), ps.add, dtype=dtype, test_nan=True)
            self.with_mode(Mode(linker="py"), ps.mul, dtype=dtype, test_nan=True)
            self.with_mode(
                Mode(linker="py"), ps.scalar_maximum, dtype=dtype, test_nan=True
            )
            self.with_mode(
                Mode(linker="py"), ps.scalar_minimum, dtype=dtype, test_nan=True
            )
            self.with_mode(
                Mode(linker="py"),
                ps.or_,
                dtype=dtype,
                test_nan=True,
                tensor_op=pt_any,
            )
            self.with_mode(
                Mode(linker="py"),
                ps.and_,
                dtype=dtype,
                test_nan=True,
                tensor_op=pt_all,
            )

    @pytest.mark.skipif(
        not pytensor.config.cxx,
        reason="G++ not available, so we need to skip this test.",
    )
    def test_c_noopt(self):
        # We need to make sure that we cover the corner cases that
        # optimizations normally cover
        self.with_mode(Mode(linker="c", optimizer=None), ps.add, dtype="floatX")

    @pytest.mark.slow
    @pytest.mark.skipif(
        not pytensor.config.cxx,
        reason="G++ not available, so we need to skip this test.",
    )
    def test_c(self):
        for dtype in ["bool", "floatX", "complex64", "complex128", "int8", "uint8"]:
            self.with_mode(Mode(linker="c"), ps.add, dtype=dtype)
            self.with_mode(Mode(linker="c"), ps.mul, dtype=dtype)
        for dtype in ["bool", "floatX", "int8", "uint8"]:
            self.with_mode(Mode(linker="c"), ps.scalar_minimum, dtype=dtype)
            self.with_mode(Mode(linker="c"), ps.scalar_maximum, dtype=dtype)
            self.with_mode(Mode(linker="c"), ps.and_, dtype=dtype, tensor_op=pt_all)
            self.with_mode(Mode(linker="c"), ps.or_, dtype=dtype, tensor_op=pt_any)
        for dtype in ["bool", "int8", "uint8"]:
            self.with_mode(Mode(linker="c"), ps.or_, dtype=dtype)
            self.with_mode(Mode(linker="c"), ps.and_, dtype=dtype)
            self.with_mode(Mode(linker="c"), ps.xor, dtype=dtype)

    @pytest.mark.slow
    @pytest.mark.skipif(
        not pytensor.config.cxx,
        reason="G++ not available, so we need to skip this test.",
    )
    def test_c_nan(self):
        for dtype in ["floatX", "complex64", "complex128"]:
            self.with_mode(Mode(linker="c"), ps.add, dtype=dtype, test_nan=True)
            self.with_mode(Mode(linker="c"), ps.mul, dtype=dtype, test_nan=True)
        for dtype in ["floatX"]:
            self.with_mode(
                Mode(linker="c"), ps.scalar_minimum, dtype=dtype, test_nan=True
            )
            self.with_mode(
                Mode(linker="c"), ps.scalar_maximum, dtype=dtype, test_nan=True
            )

    def test_infer_shape(self, dtype=None, pre_scalar_op=None):
        if dtype is None:
            dtype = pytensor.config.floatX
        for xsh, tosum in self.cases:
            x = self.type(
                dtype, shape=tuple(entry if entry == 1 else None for entry in xsh)
            )("x")
            if pre_scalar_op is not None:
                x = pre_scalar_op(x)
            if tosum is None:
                tosum = list(range(len(xsh)))
            xv = np.asarray(np.random.random(xsh), dtype=dtype)
            d = {}
            if pre_scalar_op is not None:
                xv = x.eval({x.owner.inputs[0]: xv})
                d = {pre_scalar_op: pre_scalar_op}
            self._compile_and_check(
                [x],
                [self.op(ps.add, axis=tosum, *d)(x)],
                [xv],
                self.op,
                ["local_cut_useless_reduce"],
                warn=0 not in xsh,
            )

    def test_str(self):
        op = CAReduce(ps.add, axis=None)
        assert str(op) == "CAReduce{add, axes=None}"
        op = CAReduce(ps.add, axis=(1,))
        assert str(op) == "CAReduce{add, axis=1}"

    def test_repeated_axis(self):
        x = vector("x")
        with pytest.raises(ValueError, match="repeated axis"):
            self.op(ps.add, axis=(0, 0))(x)

    def test_scalar_input(self):
        x = scalar("x")

        assert self.op(ps.add, axis=(-1,))(x).eval({x: 5}) == 5

        with pytest.raises(
            np.exceptions.AxisError,
            match=re.escape("axis (-2,) is out of bounds for array of dimension 0"),
        ):
            self.op(ps.add, axis=(-2,))(x)


class TestBitOpReduceGrad:
    def setup_method(self):
        self.rng = np.random.default_rng(unittest_tools.fetch_seed())

    def test_all_grad(self):
        x = bmatrix("x")
        x_all = x.all()
        gx = pytensor.grad(x_all, x)
        f = pytensor.function([x], gx)
        x_random = self.rng.binomial(n=1, p=0.5, size=(5, 7)).astype("int8")
        for x_val in (x_random, np.zeros_like(x_random), np.ones_like(x_random)):
            gx_val = f(x_val)
            assert gx_val.shape == x_val.shape
            assert np.all(gx_val == 0)

    def test_any_grad(self):
        x = bmatrix("x")
        x_all = x.any()
        gx = pytensor.grad(x_all, x)
        f = pytensor.function([x], gx)
        x_random = self.rng.binomial(n=1, p=0.5, size=(5, 7)).astype("int8")
        for x_val in (x_random, np.zeros_like(x_random), np.ones_like(x_random)):
            gx_val = f(x_val)
            assert gx_val.shape == x_val.shape
            assert np.all(gx_val == 0)


def check_elemwise_runtime_broadcast(mode):
    """Check we emmit a clear error when runtime broadcasting would occur according to Numpy rules."""
    x_v = matrix("x")
    m_v = vector("m")

    z_v = x_v - m_v
    f = pytensor.function([x_v, m_v], z_v, mode=mode)

    # Test invalid broadcasting by either x or m
    for x_sh, m_sh in [((2, 1), (3,)), ((2, 3), (1,))]:
        x = np.ones(x_sh).astype(config.floatX)
        m = np.zeros(m_sh).astype(config.floatX)

        # This error is introduced by PyTensor, so it's the same across different backends
        with pytest.raises(ValueError, match="Runtime broadcasting not allowed"):
            f(x, m)

    x = np.ones((2, 3)).astype(config.floatX)
    m = np.zeros((1,)).astype(config.floatX)

    x = np.ones((2, 4)).astype(config.floatX)
    m = np.zeros((3,)).astype(config.floatX)
    # This error is backend specific, and may have different types
    with pytest.raises((ValueError, TypeError)):
        f(x, m)


class TestElemwise(unittest_tools.InferShapeTester):
    def test_elemwise_grad_bool(self):
        x = scalar(dtype="bool")
        y = bscalar()
        z = x * y
        dx, dy = pytensor.grad(z, [x, y])

    def test_infer_shape(self):
        for s_left, s_right in [
            ((5, 6), (5, 6)),
            ((5, 6), (5, 1)),
            ((5, 6), (1, 6)),
            ((5, 1), (5, 6)),
            ((1, 6), (5, 6)),
            ((2, 3, 4, 5), (2, 3, 4, 5)),
            ((2, 3, 4, 5), (2, 3, 1, 5)),
            ((2, 3, 4, 5), (1, 3, 4, 5)),
            ((2, 1, 4, 5), (2, 3, 4, 5)),
            ((2, 3, 4, 1), (2, 3, 4, 5)),
        ]:
            dtype = pytensor.config.floatX
            t_left = TensorType(
                dtype, shape=tuple(entry if entry == 1 else None for entry in s_left)
            )()
            t_right = TensorType(
                dtype, shape=tuple(entry if entry == 1 else None for entry in s_right)
            )()
            t_left_val = np.zeros(s_left, dtype=dtype)
            t_right_val = np.zeros(s_right, dtype=dtype)
            self._compile_and_check(
                [t_left, t_right],
                [Elemwise(ps.add)(t_left, t_right)],
                [t_left_val, t_right_val],
                Elemwise,
            )

    def test_input_dimensions_overflow(self):
        # Elemwise.perform used to compute the product
        # of input shapes to check if there was a zero in them,
        # it overflowed in this case.
        a, b, c, d, e, f = vectors("abcdef")
        s = a + b + c + d + e + f
        g = pytensor.function([a, b, c, d, e, f], s, mode=Mode(linker="py"))
        g(*[np.zeros(2**11, config.floatX) for i in range(6)])

    def test_runtime_broadcast_python(self):
        check_elemwise_runtime_broadcast(Mode(linker="py"))

    @pytest.mark.skipif(
        not pytensor.config.cxx,
        reason="G++ not available, so we need to skip this test.",
    )
    def test_runtime_broadcast_c(self):
        check_elemwise_runtime_broadcast(Mode(linker="c"))

    def test_str(self):
        op = Elemwise(ps.add, inplace_pattern={0: 0}, name=None)
        assert str(op) == "Add"
        op = Elemwise(ps.add, inplace_pattern=None, name="my_op")
        assert str(op) == "my_op"

    def test_partial_static_shape_info(self):
        """Make sure that `Elemwise.infer_shape` can handle changes in the static shape information during rewriting."""

        x = TensorType("floatX", shape=(None, None))()
        z = Elemwise(ps.add)(x, x)

        x_inferred_shape = (ps.constant(1), ps.constant(1))

        res_shape = z.owner.op.infer_shape(
            None, z.owner, [x_inferred_shape, x_inferred_shape]
        )

        assert len(res_shape) == 1
        assert len(res_shape[0]) == 2
        assert get_scalar_constant_value(res_shape[0][0]) == 1
        assert get_scalar_constant_value(res_shape[0][1]) == 1

    def test_infer_shape_multi_output(self):
        class CustomElemwise(Elemwise):
            def make_node(self, *args):
                res = super().make_node(*args)
                return Apply(
                    self,
                    res.inputs,
                    # Return two outputs
                    [
                        TensorType(dtype="float64", shape=(None, None))()
                        for i in range(2)
                    ],
                )

        custom_elemwise = CustomElemwise(ps.add)

        z_1, z_2 = custom_elemwise(
            as_tensor_variable(np.eye(1)),
            as_tensor_variable(np.eye(1)),
        )
        in_1_shape = (ps.constant(1), ps.constant(1))
        outs = z_1.owner.op.infer_shape(None, z_1.owner, [in_1_shape, in_1_shape])
        for out in outs:
            assert out[0].eval() == 1
            assert out[1].eval() == 1

        z_1, z_2 = custom_elemwise(
            as_tensor_variable(np.eye(1)), as_tensor_variable(np.eye(3))
        )
        in_2_shape = (ps.constant(3), ps.constant(3))
        outs = z_1.owner.op.infer_shape(None, z_1.owner, [in_1_shape, in_2_shape])
        for out in outs:
            assert out[0].eval() == 3
            assert out[1].eval() == 3

    def test_shape_types(self):
        x = tensor(dtype=np.float64, shape=(None, 1))
        y = tensor(dtype=np.float64, shape=(50, 10))

        z = x * y

        assert isinstance(z.owner.op, Elemwise)

        (out_shape,) = z.owner.op.infer_shape(None, z.owner, [(lscalar(), 1), (50, 10)])

        assert all(isinstance(v.type, TensorType) for v in out_shape)

    def test_static_shape_unary(self):
        x = tensor(dtype="float64", shape=(None, 0, 1, 5))
        assert exp(x).type.shape == (None, 0, 1, 5)

    def test_static_shape_binary(self):
        x = tensor(dtype="float64", shape=(None, 5))
        y = tensor(dtype="float64", shape=(None, 5))
        assert (x + y).type.shape == (None, 5)

        x = tensor(dtype="float64", shape=(None, 5))
        y = tensor(dtype="float64", shape=(10, 5))
        assert (x + y).type.shape == (10, 5)

        x = tensor(dtype="float64", shape=(1, 5))
        y = tensor(dtype="float64", shape=(10, 5))
        assert (x + y).type.shape == (10, 5)

        x = tensor(dtype="float64", shape=(None, 1))
        y = tensor(dtype="float64", shape=(1, 1))
        assert (x + y).type.shape == (None, 1)

        x = tensor(dtype="float64", shape=(0, 0, 0))
        y = tensor(dtype="float64", shape=(0, 1, None))
        assert (x + y).type.shape == (0, 0, 0)

    def test_invalid_static_shape(self):
        x = tensor(dtype="float64", shape=(2,))
        y = tensor(dtype="float64", shape=(3,))
        with pytest.raises(
            ValueError,
            match=re.escape("Incompatible Elemwise input shapes [(2,), (3,)]"),
        ):
            x + y

    @pytest.mark.parametrize(
        "shape_x, shape_y, op, np_op",
        [
            ((3, 5), (7, 1, 3), pt.add, np.add),
            ((2, 3), (1, 4), pt.mul, np.multiply),
        ],
    )
    def test_outer(self, shape_x, shape_y, op, np_op):
        x = tensor(dtype=np.float64, shape=shape_x)
        y = tensor(dtype=np.float64, shape=shape_y)

        z = op.outer(x, y)

        f = function([x, y], z)
        x1 = np.ones(shape_x)
        y1 = np.ones(shape_y)

        np.testing.assert_array_equal(f(x1, y1), np_op.outer(x1, y1))


def test_not_implemented_elemwise_grad():
    # Regression test for unimplemented gradient in an Elemwise Op.

    class TestOp(ps.ScalarOp):
        def __init__(self):
            self.output_types_preference = ps.upgrade_to_float

        def impl(self, n, x):
            return x * n

        def grad(self, inputs, gout):
            (n, x) = inputs
            (gz,) = gout
            dy_dx = n
            return [pytensor.gradient.grad_not_implemented(self, 0, n), gz * dy_dx]

    test_op = Elemwise(TestOp())
    x = scalar()
    assert isinstance(pytensor.gradient.grad(test_op(2, x), x), Variable)

    # Verify that trying to use the not implemented gradient fails.
    with pytest.raises(pytensor.gradient.NullTypeGradError):
        pytensor.gradient.grad(test_op(x, 2), x)


class TestVectorize:
    def test_elemwise(self):
        vec = tensor(shape=(None,))
        mat = tensor(shape=(None, None))

        node = exp(vec).owner
        vect_node = vectorize_node(node, mat)
        assert vect_node.op == exp
        assert vect_node.inputs[0] is mat

    def test_dimshuffle(self):
        vec = tensor(shape=(None,))
        mat = tensor(shape=(None, None))

        node = exp(vec).owner
        vect_node = vectorize_node(node, mat)
        assert vect_node.op == exp
        assert vect_node.inputs[0] is mat

        col_mat = tensor(shape=(None, 1))
        tcol_mat = tensor(shape=(None, None, 1))
        node = col_mat.dimshuffle(0).owner  # drop column
        vect_node = vectorize_node(node, tcol_mat)
        assert isinstance(vect_node.op, DimShuffle)
        assert vect_node.op.new_order == (0, 1)
        assert vect_node.inputs[0] is tcol_mat
        assert vect_node.outputs[0].type.shape == (None, None)

    def test_CAReduce(self):
        mat = tensor(shape=(None, None))
        tns = tensor(shape=(None, None, None))

        node = pt_sum(mat).owner
        vect_node = vectorize_node(node, tns)
        assert isinstance(vect_node.op, Sum)
        assert vect_node.op.axis == (1, 2)
        assert vect_node.inputs[0] is tns

        bool_mat = tensor(dtype="bool", shape=(None, None))
        bool_tns = tensor(dtype="bool", shape=(None, None, None))
        node = pt_any(bool_mat, axis=-2).owner
        vect_node = vectorize_node(node, bool_tns)
        assert isinstance(vect_node.op, Any)
        assert vect_node.op.axis == (1,)
        assert vect_node.inputs[0] is bool_tns


def careduce_benchmark_tester(axis, c_contiguous, mode, benchmark):
    N = 256
    x_test = np.random.uniform(size=(N, N, N))
    transpose_axis = (0, 1, 2) if c_contiguous else (2, 0, 1)

    x = pytensor.shared(x_test, name="x", shape=x_test.shape)
    out = x.transpose(transpose_axis).sum(axis=axis)
    fn = pytensor.function([], out, mode=mode)

    np.testing.assert_allclose(
        fn(),
        x_test.transpose(transpose_axis).sum(axis=axis),
    )
    benchmark(fn)


@pytest.mark.parametrize(
    "axis",
    (0, 1, 2, (0, 1), (0, 2), (1, 2), None),
    ids=lambda x: f"axis={x}",
)
@pytest.mark.parametrize(
    "c_contiguous",
    (True, False),
    ids=lambda x: f"c_contiguous={x}",
)
def test_c_careduce_benchmark(axis, c_contiguous, benchmark):
    return careduce_benchmark_tester(
        axis, c_contiguous, mode="FAST_RUN", benchmark=benchmark
    )
