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, config, grad
from pytensor.compile.function import function
from pytensor.compile.mode import Mode, get_default_mode
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.link.numba import NumbaLinker
from pytensor.scalar import ScalarOp, float32, float64, int32, int64
from pytensor.scalar import add as scalar_add
from pytensor.scalar import exp as scalar_exp
from pytensor.scalar import xor as scalar_xor
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,
    dmatrix,
    lscalar,
    matrix,
    scalar,
    tensor,
    tensor3,
    vector,
    vectors,
)
from tests import unittest_tools
from tests.link.test_link import make_function
from tests.tensor.utils import (
    _bad_runtime_broadcast_binary_normal,
    inplace_func,
    integers,
    integers_uint16,
    integers_uint32,
    makeBroadcastTester,
    random,
    random_complex,
)


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=())()
        with pytest.raises(
            ValueError,
            match="maximum supported dimension for a TensorType is currently 64, found 65",
        ):
            x.dimshuffle(("x",) * 65)

    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,))
        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[-1] == 0
            # Confirm the broadcasted value in the output
            assert np.array_equiv(outputs[0].storage[0], 2039)

    @pytest.mark.skipif(
        isinstance(get_default_mode().linker, NumbaLinker),
        reason="Running with numba linker, c backend (should be) covered in another CI",
    )
    @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()
        burn_in = 1
        assert np.allclose(np.mean(block_diffs[burn_in:]), 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 = 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, {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, {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")
            op1 = op(ps.second, {0: 0})
            op2 = op(ps.second, {0: 0})
            assert op1 == op2
            e = op(ps.Second(), {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.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.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.maximum, ps.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.maximum, dtype=dtype)
            self.with_mode(Mode(linker="py"), ps.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.maximum, dtype=dtype, test_nan=True)
            self.with_mode(Mode(linker="py"), ps.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.",
    )
    @pytest.mark.skipif(
        isinstance(get_default_mode().linker, NumbaLinker),
        reason="Running with numba linker, c backend (should be) covered in another CI",
    )
    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.",
    )
    @pytest.mark.skipif(
        isinstance(get_default_mode().linker, NumbaLinker),
        reason="Running with numba linker, c backend (should be) covered in another CI",
    )
    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.minimum, dtype=dtype)
            self.with_mode(Mode(linker="c"), ps.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.",
    )
    @pytest.mark.skipif(
        isinstance(get_default_mode().linker, NumbaLinker),
        reason="Running with numba linker, c backend (should be) covered in another CI",
    )
    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.minimum, dtype=dtype, test_nan=True)
            self.with_mode(Mode(linker="c"), ps.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):
        c_mode = Mode(linker="cvm")
        check_elemwise_runtime_broadcast(c_mode)

        # Test C-backend specific error formatting
        x = dmatrix("x")
        y = dmatrix("y")
        fn = function([x, y], x * y, mode=c_mode)
        with pytest.raises(
            ValueError,
            match=r"Runtime broadcasting not allowed.*\(input\[0\]\.shape\[1\] = 4, input\[1\]\.shape\[1\] = 1\)",
        ):
            fn(np.zeros((5, 4)), np.zeros((5, 1)))

        with pytest.raises(
            ValueError,
            match=re.escape(
                "Input dimension mismatch: (input[0].shape[1] = 4, input[1].shape[1] = 3)"
            ),
        ):
            fn(np.zeros((5, 4)), np.zeros((5, 3)))

    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
    )


def test_gradient_mixed_discrete_output_scalar_op():
    class MixedDtypeScalarOp(ScalarOp):
        def make_node(self, *inputs):
            float_op = float64 if config.floatX == "float64" else float32
            int_op = int64 if config.floatX == "int64" else int32
            inputs = [float_op()]
            outputs = [float_op(), int_op()]
            return Apply(self, inputs, outputs)

        def perform(self, node, inputs, outputs):
            raise NotImplementedError()

        def L_op(self, inputs, outputs, output_gradients):
            return [inputs[0].ones_like() * output_gradients[0]]

    op = Elemwise(MixedDtypeScalarOp())
    x = vector("x")
    y, _ = op(x)
    np.testing.assert_array_equal(
        grad(y.sum(), x).eval({x: np.full((12,), np.nan, dtype=config.floatX)}),
        np.ones((12,), dtype=config.floatX),
        strict=True,
    )


@pytest.mark.filterwarnings("error")
def test_numpy_warning_suppressed():
    x = pt.scalar("x")
    y = pt.log(x)
    fn = pytensor.function([x], y, mode=Mode(linker="py"))
    assert fn(0) == -np.inf


rng = np.random.default_rng(18)
_good_add_inplace = dict(
    same_shapes=(random(2, 3, rng=rng), random(2, 3, rng=rng)),
    not_same_dimensions=(random(2, 2, rng=rng), random(2, rng=rng)),
    scalar=(random(2, 3, rng=rng), random(1, 1, rng=rng)),
    row=(random(2, 3, rng=rng), random(1, 3, rng=rng)),
    column=(random(2, 3, rng=rng), random(2, 1, rng=rng)),
    integers=(integers(2, 3, rng=rng), integers(2, 3, rng=rng)),
    uint32=(integers_uint32(2, 3, rng=rng), integers_uint32(2, 3, rng=rng)),
    uint16=(integers_uint16(2, 3, rng=rng), integers_uint16(2, 3, rng=rng)),
    # (float32, >int16) upcasts to float64 by default
    dtype_valid_mixup=(
        random(2, 3, rng=rng),
        integers(2, 3, rng=rng).astype(
            "int16" if config.floatX == "float32" else "int64"
        ),
    ),
    complex1=(random_complex(2, 3, rng=rng), random_complex(2, 3, rng=rng)),
    complex2=(random_complex(2, 3, rng=rng), random(2, 3, rng=rng)),
    empty=(np.asarray([], dtype=config.floatX), np.asarray([1], dtype=config.floatX)),
)
TestAddInplaceBroadcast = makeBroadcastTester(
    op=Elemwise(scalar_add, {0: 0}),
    expected=lambda x, y: x + y,
    good=_good_add_inplace,
    # Cannot inplace on first input if it doesn't match output dtype (upcast of inputs)
    bad_build=dict(dtype_invalid_mixup=_good_add_inplace["dtype_valid_mixup"][::-1]),
    bad_runtime=_bad_runtime_broadcast_binary_normal,
    inplace=True,
)


@pytest.mark.xfail(
    config.cycle_detection == "fast" and config.mode != "FAST_COMPILE",
    reason="Cycle detection is fast and mode is FAST_COMPILE",
)
def test_exp_inplace_grad_1():
    utt.verify_grad(
        Elemwise(scalar_exp, {0: 0}),
        [
            np.asarray(
                [
                    [1.5089518, 1.48439076, -4.7820262],
                    [2.04832468, 0.50791564, -1.58892269],
                ]
            )
        ],
    )


def test_XOR_inplace():
    dtype = [
        "int8",
        "int16",
        "int32",
        "int64",
    ]
    xor_inplace = Elemwise(scalar_xor, {0: 0})

    for dtype in dtype:
        x, y = vector(dtype=dtype), vector(dtype=dtype)
        l = np.asarray([0, 0, 1, 1], dtype=dtype)
        r = np.asarray([0, 1, 0, 1], dtype=dtype)
        ix = x
        ix = xor_inplace(ix, y)
        gn = inplace_func([x, y], ix)
        _ = gn(l, r)
        # test the in-place stuff
        assert np.all(l == np.asarray([0, 1, 1, 0])), l


def test_inplace_dtype_changed():
    with pytensor.config.change_flags(cast_policy="numpy+floatX", floatX="float64"):
        x = pt.vector("x", dtype="float32")
        y = pt.vector("y", dtype="int32")
        with pytensor.config.change_flags(floatX="float32"):
            out = pt.add(x, y)

        assert out.dtype == "float32"
        with pytensor.config.change_flags(floatX="float32"):
            fn32 = pytensor.function(
                [In(x, mutable=True), In(y, mutable=True)],
                out,
                mode="fast_run",
            )
        assert fn32.maker.fgraph.outputs[0].owner.op.destroy_map == {0: [0]}

        with pytensor.config.change_flags(floatX="float64"):
            fn64 = pytensor.function(
                [In(x, mutable=True), In(y, mutable=True)],
                out,
                mode="fast_run",
            )
        assert fn64.maker.fgraph.outputs[0].owner.op.destroy_map == {}
