import re
from itertools import product

import numpy as np
import pytest
import scipy.linalg

import pytensor
from pytensor import In, config, function, scan
from pytensor.compile import get_default_mode, get_mode
from pytensor.compile.function.types import add_supervisor_to_fgraph
from pytensor.gradient import grad
from pytensor.graph import Apply, FunctionGraph, Op, rewrite_graph
from pytensor.graph.replace import _vectorize_node, vectorize_graph
from pytensor.link.numba import NumbaLinker
from pytensor.raise_op import assert_op
from pytensor.tensor import (
    matrices,
    matrix,
    ones_like,
    scalar,
    tensor,
    vector,
)
from pytensor.tensor.blockwise import (
    Blockwise,
    BlockwiseWithCoreShape,
    vectorize_node_fallback,
)
from pytensor.tensor.nlinalg import MatrixInverse, eig
from pytensor.tensor.random import normal
from pytensor.tensor.random.op import default_rng
from pytensor.tensor.rewriting.blas import specialize_matmul_to_batched_dot
from pytensor.tensor.slinalg import (
    Cholesky,
    Solve,
    SolveBase,
    cho_solve,
    cholesky,
    solve,
    solve_triangular,
)
from pytensor.tensor.utils import _parse_gufunc_signature


@pytest.mark.xfail(
    condition=isinstance(get_default_mode().linker, NumbaLinker),
    raises=TypeError,
    reason="Numba scalar blockwise obj-mode fallback fails: https://github.com/pymc-devs/pytensor/issues/1760",
)
def test_perform_method_per_node():
    """Confirm that Blockwise uses one perform method per node.

    This is important if the perform method requires node information (such as dtypes)
    """

    class NodeDependentPerformOp(Op):
        def make_node(self, x):
            return Apply(self, [x], [x.type()])

        def perform(self, node, inputs, outputs):
            [x] = inputs
            if node.inputs[0].type.dtype.startswith("float"):
                y = x + 1
            else:
                y = x - 1
            outputs[0][0] = y

    blockwise_op = Blockwise(core_op=NodeDependentPerformOp(), signature="()->()")
    x = tensor("x", shape=(3,), dtype="float32")
    y = tensor("y", shape=(3,), dtype="int32")

    out_x = blockwise_op(x)
    out_y = blockwise_op(y)
    fn = pytensor.function([x, y], [out_x, out_y])
    [op1, op2] = [node.op for node in fn.maker.fgraph.apply_nodes]
    # Confirm both nodes have the same Op
    assert isinstance(op1, Blockwise | BlockwiseWithCoreShape) and isinstance(
        op1.core_op, NodeDependentPerformOp
    )
    assert isinstance(op2, Blockwise | BlockwiseWithCoreShape) and isinstance(
        op2.core_op, NodeDependentPerformOp
    )
    # assert op1 is op2  # Not true in the Numba backend

    res_out_x, res_out_y = fn(np.zeros(3, dtype="float32"), np.zeros(3, dtype="int32"))
    np.testing.assert_array_equal(res_out_x, np.ones(3, dtype="float32"))
    np.testing.assert_array_equal(res_out_y, -np.ones(3, dtype="int32"))


def test_vectorize_blockwise():
    mat = tensor(shape=(None, None))
    tns = tensor(shape=(None, None, None))

    # Something that falls back to Blockwise
    out = MatrixInverse()(mat)
    vect_node = vectorize_graph(out, {mat: tns}).owner
    assert isinstance(vect_node.op, Blockwise) and isinstance(
        vect_node.op.core_op, MatrixInverse
    )
    assert vect_node.op.signature == ("(m,m)->(m,m)")
    assert vect_node.inputs[0] is tns

    # Useless blockwise
    tns4 = tensor(shape=(5, None, None, None))
    new_vect_node = vectorize_graph(vect_node.out, {tns: tns4}).owner
    assert new_vect_node.op is vect_node.op
    assert isinstance(new_vect_node.op, Blockwise) and isinstance(
        new_vect_node.op.core_op, MatrixInverse
    )
    assert new_vect_node.inputs[0] is tns4


def test_vectorize_node_fallback_unsupported_type():
    rng = default_rng()
    node = normal(rng=rng).owner

    with pytest.raises(
        NotImplementedError,
        match=re.escape(
            'Cannot vectorize node normal_rv{"(),()->()"}('
            "DefaultGeneratorMakerOp.0, NoneConst{None}, 0.0, 1.0)"
            " with input DefaultGeneratorMakerOp.0 of type RandomGeneratorType"
        ),
    ):
        vectorize_node_fallback(node.op, node, *node.inputs)


def check_blockwise_runtime_broadcasting(mode):
    a = tensor("a", shape=(None, 3, 5))
    b = tensor("b", shape=(None, 5, 3))

    out = a @ b
    fn = function(
        [a, b],
        out,
        mode=get_mode(mode).excluding(specialize_matmul_to_batched_dot.__name__),
    )
    assert isinstance(
        fn.maker.fgraph.outputs[0].owner.op, Blockwise | BlockwiseWithCoreShape
    )

    for valid_test_values in [
        (
            np.ones((2, 3, 5)).astype(config.floatX),
            np.ones((2, 5, 3)).astype(config.floatX),
        ),
        (
            np.ones((1, 3, 5)).astype(config.floatX),
            np.ones((1, 5, 3)).astype(config.floatX),
        ),
    ]:
        batch_dim = valid_test_values[0].shape[0]
        np.testing.assert_allclose(
            fn(*valid_test_values), np.full((batch_dim, 3, 3), 5.0)
        )

    possible_err_messages = [
        "Runtime broadcasting not allowed",
        "has an incompatible shape in axis",
        "Incompatible vectorized shapes",
    ]
    err_msg = f"({'|'.join(possible_err_messages)})"
    for invalid_test_values in [
        (
            np.ones((1, 3, 5)).astype(config.floatX),
            np.ones((2, 5, 3)).astype(config.floatX),
        ),
        (
            np.ones((2, 3, 5)).astype(config.floatX),
            np.ones((1, 5, 3)).astype(config.floatX),
        ),
    ]:
        with pytest.raises(ValueError, match=err_msg):
            fn(*invalid_test_values)

    invalid_test_values = (
        np.ones((2, 3, 5)).astype(config.floatX),
        np.ones((3, 5, 3)).astype(config.floatX),
    )
    # Error message is backend specific
    with pytest.raises(ValueError):
        fn(*invalid_test_values)


@pytest.mark.parametrize("mode", ("FAST_COMPILE", "FAST_RUN"))
def test_runtime_broadcast(mode):
    check_blockwise_runtime_broadcasting(mode)


class MyTestOp(Op):
    def make_node(self, *inputs):
        return Apply(self, inputs, [i.type() for i in inputs])

    def perform(self, *args, **kwargs):
        raise NotImplementedError("Test Op should not be present in final graph")


my_test_op = MyTestOp()


def test_vectorize_node_default_signature():
    vec = tensor(shape=(None,))
    mat = tensor(shape=(5, None))
    node = my_test_op.make_node(vec, mat)

    vect_node = _vectorize_node(node.op, node, mat, mat)
    assert isinstance(vect_node.op, Blockwise) and isinstance(
        vect_node.op.core_op, MyTestOp
    )
    assert vect_node.op.signature == ("(i00),(i10,i11)->(o00),(o10,o11)")

    with pytest.raises(
        ValueError, match="Signature not provided nor found in core_op MyTestOp"
    ):
        Blockwise(my_test_op)

    vect_node = Blockwise(my_test_op, signature="(m),(n)->(m),(n)").make_node(vec, mat)
    assert vect_node.outputs[0].type.shape == (
        5,
        None,
    )
    assert vect_node.outputs[0].type.shape == (
        5,
        None,
    )


def test_blockwise_shape():
    # Single output
    inp = tensor(shape=(5, None, None))
    inp_test = np.zeros((5, 4, 3), dtype=config.floatX)

    # Shape can be inferred from inputs
    op = Blockwise(my_test_op, signature="(m, n) -> (n, m)")
    out = op(inp)
    assert out.type.shape == (5, None, None)

    shape_fn = pytensor.function([inp], out.shape)
    assert not any(
        isinstance(getattr(n.op, "core_op", n.op), MyTestOp)
        for n in shape_fn.maker.fgraph.apply_nodes
    )
    assert tuple(shape_fn(inp_test)) == (5, 3, 4)

    # Shape can only be partially inferred from inputs
    op = Blockwise(my_test_op, signature="(m, n) -> (m, k)")
    out = op(inp)
    assert out.type.shape == (5, None, None)

    shape_fn = pytensor.function([inp], out.shape)
    assert any(
        isinstance(getattr(n.op, "core_op", n.op), MyTestOp)
        for n in shape_fn.maker.fgraph.apply_nodes
    )

    shape_fn = pytensor.function([inp], out.shape[:-1])
    assert not any(
        isinstance(getattr(n.op, "core_op", n.op), MyTestOp)
        for n in shape_fn.maker.fgraph.apply_nodes
    )
    assert tuple(shape_fn(inp_test)) == (5, 4)

    # Mutiple outputs
    inp1 = tensor(shape=(7, 1, None, None))
    inp2 = tensor(shape=(1, 5, None, None))
    inp1_test = np.zeros((7, 1, 4, 3), dtype=config.floatX)
    inp2_test = np.zeros((1, 5, 4, 3), dtype=config.floatX)

    op = Blockwise(my_test_op, signature="(m, n), (m, n) -> (n, m), (m, k)")
    outs = op(inp1, inp2)
    assert outs[0].type.shape == (7, 5, None, None)
    assert outs[1].type.shape == (7, 5, None, None)

    shape_fn = pytensor.function([inp1, inp2], [out.shape for out in outs])
    assert any(
        isinstance(getattr(n.op, "core_op", n.op), MyTestOp)
        for n in shape_fn.maker.fgraph.apply_nodes
    )

    shape_fn = pytensor.function([inp1, inp2], outs[0].shape)
    assert not any(
        isinstance(getattr(n.op, "core_op", n.op), MyTestOp)
        for n in shape_fn.maker.fgraph.apply_nodes
    )
    assert tuple(shape_fn(inp1_test, inp2_test)) == (7, 5, 3, 4)

    shape_fn = pytensor.function([inp1, inp2], [outs[0].shape, outs[1].shape[:-1]])
    assert not any(
        isinstance(getattr(n.op, "core_op", n.op), MyTestOp)
        for n in shape_fn.maker.fgraph.apply_nodes
    )
    assert tuple(shape_fn(inp1_test, inp2_test)[0]) == (7, 5, 3, 4)
    assert tuple(shape_fn(inp1_test, inp2_test)[1]) == (7, 5, 4)


def test_blockwise_infer_core_shape():
    class TestOpWithInferShape(Op):
        def make_node(self, a, b):
            assert a.type.ndim == 1
            assert b.type.ndim == 1
            # Simulate make_node that introduces operations on inputs
            a_identity = a.copy()
            b_identity = b.copy()

            c = tensor(shape=(None,))
            d = tensor(shape=(None,))
            return Apply(self, [a_identity, b_identity], [c, d])

        def perform(self, node, inputs, outputs):
            a, b = inputs
            c, d = outputs
            c[0] = np.arange(a.size + b.size, dtype=config.floatX)
            d[0] = np.arange(a.sum() + b.sum(), dtype=config.floatX)

        def infer_shape(self, fgraph, node, input_shapes):
            # First output shape depends only on input_shapes
            # Second output shape depends on input values
            a_identity, b_identity = node.inputs
            # Simulate shape depending on original inputs, not the ones that go directly into the node
            a = a_identity.owner.inputs[0]
            b = b_identity.owner.inputs[0]
            [(a_shape,), (b_shape,)] = input_shapes
            return (a_shape + b_shape,), (a.sum() + b.sum(),)

    blockwise_op = Blockwise(
        core_op=TestOpWithInferShape(), signature="(a),(b)->(c),(d)"
    )

    a = tensor("a", shape=(5, 3))
    b = tensor("b", shape=(1, 4))
    c, d = blockwise_op(a, b)
    assert c.type.shape == (5, None)
    assert d.type.shape == (5, None)

    c_shape_fn = pytensor.function([a, b], c.shape)
    # c_shape can be computed from the input shapes alone
    assert not any(
        isinstance(getattr(n.op, "core_op", n.op), TestOpWithInferShape)
        for n in c_shape_fn.maker.fgraph.apply_nodes
    )

    d_shape_fn = pytensor.function([a, b], d.shape)
    # d_shape cannot be computed from the input shapes alone
    assert any(
        isinstance(getattr(n.op, "core_op", n.op), TestOpWithInferShape)
        for n in d_shape_fn.maker.fgraph.apply_nodes
    )

    a_test = np.zeros(a.type.shape, dtype=a.type.dtype)
    b_test = np.zeros(b.type.shape, dtype=b.type.dtype)
    assert tuple(c_shape_fn(a_test, b_test)) == (5, 7)
    assert tuple(d_shape_fn(a_test, b_test)) == (5, 0)


def test_infer_shape_literals():
    # Define a CoreOp whose infer_shape is (symbolic operation on itself, literal)
    # Then tell Blockwise that the first dimension is constant.
    # The Op has no perform method, so it will fail to evaluate if the infer_shape of that dimension isn't ignored
    class TestCoreOp(Op):
        def make_node(self, x):
            assert x.type.ndim == 0
            return Apply(self, [x], [matrix()])

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

        def infer_shape(self, fgraph, node, input_shapes):
            y = node.outputs[0]
            # Apparently it's valid to return integers in infer_shape.
            # DimShuffle does this. Modify test if that is no longer allowed.
            return [(y[0][0].astype(int), 3)]

    op = Blockwise(TestCoreOp(), signature="()->(2,a)")
    x = scalar("x")
    y = op(x)

    fn = function([x], y.shape)
    assert tuple(fn(0)) == (2, 3)


class BlockwiseOpTester:
    """Base class to test Blockwise works for specific Ops"""

    core_op = None
    signature = None
    batcheable_axes = None

    @classmethod
    def setup_class(cls):
        seed = sum(map(ord, str(cls.core_op) + cls.signature))
        cls.rng = np.random.default_rng(seed)
        cls.params_sig, cls.outputs_sig = _parse_gufunc_signature(cls.signature)
        if cls.batcheable_axes is None:
            cls.batcheable_axes = list(range(len(cls.params_sig)))
        batch_shapes = [(), (1,), (5,), (1, 1), (1, 5), (3, 1), (3, 5)]
        cls.test_batch_shapes = list(
            product(batch_shapes, repeat=len(cls.batcheable_axes))
        )
        cls.block_op = Blockwise(core_op=cls.core_op, signature=cls.signature)

    @staticmethod
    def parse_shape(shape: tuple[str | int, ...]) -> tuple[int, ...]:
        """
        Convert (5, "m", "n") -> (5, 7, 11)
        """
        mapping = {"m": 7, "n": 11, "k": 19}
        return tuple(mapping.get(p, p) for p in shape)

    def create_testvals(self, shape):
        return self.rng.normal(size=self.parse_shape(shape)).astype(config.floatX)

    def create_batched_inputs(self, batch_idx: int | None = None):
        for batch_shapes in self.test_batch_shapes:
            vec_inputs = []
            vec_inputs_testvals = []
            for idx, (batch_shape, param_sig) in enumerate(
                zip(batch_shapes, self.params_sig, strict=True)
            ):
                if batch_idx is not None and idx != batch_idx:
                    # Skip out combinations in which other inputs are batched
                    if batch_shape != ():
                        break
                vec_inputs.append(tensor(shape=batch_shape + (None,) * len(param_sig)))
                vec_inputs_testvals.append(
                    self.create_testvals(shape=batch_shape + param_sig)
                )
            else:  # no-break
                yield vec_inputs, vec_inputs_testvals

    def test_perform(self):
        base_inputs = [
            tensor(shape=(None,) * len(param_sig)) for param_sig in self.params_sig
        ]
        core_func = pytensor.function(base_inputs, self.core_op(*base_inputs))

        def inp_copy_core_func(*args):
            # Work-around for https://github.com/numba/numba/issues/10357
            # numpy vectorize passes non-writeable arrays to the inner function
            return core_func(*(arg.copy() for arg in args))

        np_func = np.vectorize(inp_copy_core_func, signature=self.signature)

        for vec_inputs, vec_inputs_testvals in self.create_batched_inputs():
            pt_func = pytensor.function(vec_inputs, self.block_op(*vec_inputs))
            if len(self.outputs_sig) != 1:
                raise NotImplementedError("Did not implement test for multi-output Ops")
            np.testing.assert_allclose(
                pt_func(*vec_inputs_testvals),
                np_func(*vec_inputs_testvals),
                rtol=1e-7 if config.floatX == "float64" else 1e-5,
                atol=1e-7 if config.floatX == "float64" else 1e-5,
            )

    def test_grad(self):
        base_inputs = [
            tensor(shape=(None,) * len(param_sig)) for param_sig in self.params_sig
        ]
        out = self.core_op(*base_inputs).sum()
        # Create separate numpy vectorized functions for each input

        def copy_inputs_wrapper(fn):
            # Work-around for https://github.com/numba/numba/issues/10357
            # numpy vectorize passes non-writeable arrays to the inner function
            def copy_fn(*args):
                return fn(*(arg.copy() for arg in args))

            return copy_fn

        np_funcs = []
        for i, inp in enumerate(base_inputs):
            core_grad_func = pytensor.function(base_inputs, grad(out, wrt=inp))
            params_sig = self.signature.split("->")[0]
            param_sig = f"({','.join(self.params_sig[i])})"
            grad_sig = f"{params_sig}->{param_sig}"

            np_func = np.vectorize(
                copy_inputs_wrapper(core_grad_func),
                signature=grad_sig,
            )
            np_funcs.append(np_func)

        # We test gradient wrt to one batched input at a time
        for test_input_idx in range(len(base_inputs)):
            for vec_inputs, vec_inputs_testvals in self.create_batched_inputs(
                batch_idx=test_input_idx
            ):
                out = self.block_op(*vec_inputs).sum()
                pt_func = pytensor.function(
                    vec_inputs, grad(out, wrt=vec_inputs[test_input_idx])
                )
                pt_out = pt_func(*vec_inputs_testvals)
                np_out = np_funcs[test_input_idx](*vec_inputs_testvals)
                np.testing.assert_allclose(
                    pt_out,
                    np_out,
                    rtol=1e-7 if config.floatX == "float64" else 1e-5,
                    atol=1e-6 if config.floatX == "float64" else 2e-2,
                )


class MatrixOpBlockwiseTester(BlockwiseOpTester):
    def create_testvals(self, shape):
        # Return a posdef matrix
        X = super().create_testvals(shape)
        return np.einsum("...ij,...kj->...ik", X, X)


class TestCholesky(MatrixOpBlockwiseTester):
    core_op = Cholesky(lower=True)
    signature = "(m, m) -> (m, m)"


class TestMatrixInverse(MatrixOpBlockwiseTester):
    core_op = MatrixInverse()
    signature = "(m, m) -> (m, m)"


class TestSolveVector(BlockwiseOpTester):
    core_op = Solve(lower=True, b_ndim=1)
    signature = "(m, m),(m) -> (m)"


class TestSolveMatrix(BlockwiseOpTester):
    core_op = Solve(lower=True, b_ndim=2)
    signature = "(m, m),(m, n) -> (m, n)"


def test_cop_with_params():
    matrix_assert = Blockwise(core_op=assert_op, signature="(x1,x2),()->(x1,x2)")

    x = tensor("x", shape=(5, None, None), dtype="float64")
    x_shape = matrix_assert(x, (x >= 0).all())

    fn = pytensor.function([x], x_shape)
    [fn_out] = fn.maker.fgraph.outputs
    op = fn_out.owner.op
    assert (
        isinstance(op, Blockwise | BlockwiseWithCoreShape) and op.core_op == assert_op
    ), "Blockwise should be in final graph"

    np.testing.assert_allclose(
        fn(np.zeros((5, 3, 2))),
        np.zeros((5, 3, 2)),
    )

    with pytest.raises(AssertionError):
        fn(np.zeros((5, 3, 2)) - 1)


@pytest.mark.skipif(
    config.mode == "FAST_COMPILE",
    reason="inplace rewrites disabled when mode is FAST_COMPILE",
)
class TestInplace:
    @pytest.mark.parametrize("is_batched", (False, True))
    def test_cholesky(self, is_batched):
        X = tensor("X", shape=(5, None, None) if is_batched else (None, None))
        L = cholesky(X, lower=True)
        f = function([In(X, mutable=True)], L)

        assert not L.owner.op.core_op.destroy_map

        if is_batched:
            [cholesky_op] = [
                node.op.core_op
                for node in f.maker.fgraph.apply_nodes
                if isinstance(node.op, Blockwise | BlockwiseWithCoreShape)
                and isinstance(node.op.core_op, Cholesky)
            ]
        else:
            [cholesky_op] = [
                node.op
                for node in f.maker.fgraph.apply_nodes
                if isinstance(node.op, Cholesky)
            ]
        assert cholesky_op.destroy_map == {0: [0]}

        rng = np.random.default_rng(441 + is_batched)
        X_val = rng.normal(size=(10, 10)).astype(config.floatX)
        X_val_in = X_val @ X_val.T
        if is_batched:
            X_val_in = np.broadcast_to(X_val_in, (5, *X_val_in.shape)).copy()
        X_val_in_copy = X_val_in.copy()

        f(X_val_in)

        np.testing.assert_allclose(
            X_val_in,
            np.linalg.cholesky(X_val_in_copy),
            atol=1e-5 if config.floatX == "float32" else 0,
        )

    @pytest.mark.parametrize("batched_A", (False, True))
    @pytest.mark.parametrize("batched_b", (False, True))
    @pytest.mark.parametrize("solve_fn", (solve, solve_triangular, cho_solve))
    def test_solve(self, solve_fn, batched_A, batched_b):
        A = tensor("A", shape=(5, 3, 3) if batched_A else (3, 3))
        b = tensor("b", shape=(5, 3) if batched_b else (3,))
        if solve_fn == cho_solve:
            # Special signature for cho_solve
            x = solve_fn((A, True), b, b_ndim=1)
        else:
            x = solve_fn(A, b, b_ndim=1)

        mode = get_default_mode().excluding(
            "batched_vector_b_solve_to_matrix_b_solve",
            "reuse_decomposition_multiple_solves",
        )
        fn = function([In(A, mutable=True), In(b, mutable=True)], x, mode=mode)

        op = fn.maker.fgraph.outputs[0].owner.op
        if batched_A or batched_b:
            assert isinstance(op, Blockwise | BlockwiseWithCoreShape) and isinstance(
                op.core_op, SolveBase
            )
            if batched_A and not batched_b:
                if solve_fn == solve:
                    assert op.destroy_map == {0: [0]}
                else:
                    # SolveTriangular does not destroy A
                    assert op.destroy_map == {}
            else:
                assert op.destroy_map == {0: [1]}
        else:
            assert isinstance(op, SolveBase)
            assert op.destroy_map == {0: [1]}

        # We test with an F_CONTIGUOUS (core) A as only that will be destroyed by scipy
        rng = np.random.default_rng(
            487 + batched_A + 2 * batched_b + sum(map(ord, solve_fn.__name__))
        )
        A_val = np.swapaxes(rng.normal(size=A.type.shape).astype(A.type.dtype), -1, -2)
        b_val = np.random.normal(size=b.type.shape).astype(b.type.dtype)
        A_val_copy = A_val.copy()
        b_val_copy = b_val.copy()
        out = fn(A_val, b_val)

        if solve_fn == cho_solve:

            def core_scipy_fn(A, b):
                return scipy.linalg.cho_solve((A, True), b)

        else:
            core_scipy_fn = getattr(scipy.linalg, solve_fn.__name__)
        expected_out = np.vectorize(core_scipy_fn, signature="(m,m),(m)->(m)")(
            A_val_copy, b_val_copy
        )
        np.testing.assert_allclose(
            out, expected_out, atol=1e-4 if config.floatX == "float32" else 0
        )

        # Confirm input was destroyed
        assert (A_val == A_val_copy).all() == (op.destroy_map.get(0, None) != [0])
        assert (b_val == b_val_copy).all() == (op.destroy_map.get(0, None) != [1])


def test_gradient_mixed_discrete_output_core_op():
    class MixedDtypeCoreOp(Op):
        gufunc_signature = "()->(),()"
        itypes = [scalar().type]
        otypes = [scalar().type, scalar(dtype=int).type]

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

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

    op = Blockwise(MixedDtypeCoreOp())
    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,
    )


def test_blockwise_grad_core_type():
    class StrictCoreTypeOp(Op):
        def make_node(self, x):
            assert x.type.shape[-1] == 2
            return Apply(self, [x], [x.type()])

        def perform(self, node, inputs, output_storage):
            output_storage[0][0] = inputs[0] + 1

        def L_op(self, inputs, outputs, output_grads):
            [x] = inputs
            assert x.type.shape == (2,)
            return [x.zeros_like()]

    strict_core_type_op = StrictCoreTypeOp()
    block_strict_core_type_op = Blockwise(strict_core_type_op, signature="(a)->(a)")

    x = tensor("x", shape=(5, 2), dtype="float64")
    y = block_strict_core_type_op(x)
    assert y.type.shape == (5, 2)

    grad_y = grad(y.sum(), x)
    assert grad_y.type.shape == (5, 2)
    np.testing.assert_allclose(
        grad_y.eval({x: np.ones((5, 2))}),
        np.zeros((5, 2)),
    )


def test_scan_gradient_core_type():
    n_steps = 3
    seq = tensor("seq", shape=(n_steps, 1), dtype="float64")
    out = scan(
        lambda s: s,
        sequences=[seq],
        n_steps=n_steps,
        return_updates=False,
    )

    vec_seq = tensor("vec_seq", shape=(None, n_steps, 1), dtype="float64")
    vec_out = vectorize_graph(out, replace={seq: vec_seq})
    grad_sit_sot0 = grad(vec_out.sum(), vec_seq)

    np.testing.assert_allclose(
        grad_sit_sot0.eval({vec_seq: np.ones((4, n_steps, 1))}),
        np.ones((4, n_steps, 1)),
    )


def test_partial_inplace():
    class CoreOp(Op):
        __props__ = ("inplace",)

        def __init__(self, inplace):
            self.inplace = tuple(inplace)
            self.destroy_map = {i: [i] for i in inplace}

        def inplace_on_inputs(self, allowed_inplace_inputs):
            return type(self)(inplace=allowed_inplace_inputs)

        def make_node(self, x, y, z):
            return Apply(self, [x, y, z], [x.type(), y.type(), z.type()])

        def perform(self, node, inputs, outputs):
            [x, y, z] = inputs
            if 0 not in self.inplace:
                x = x.copy()
            if 1 not in self.inplace:
                y = y.copy()
            if 2 not in self.inplace:
                z = z.copy()
            outputs[0][0] = x
            outputs[1][0] = y
            outputs[2][0] = z

    core_op = CoreOp(inplace=())
    blockwise_op = Blockwise(core_op, signature="(),(),()->(),(),()")
    x, y, z = matrices("xyz")

    # All can be inplaced
    out = blockwise_op(x.T, y.T, z.T)
    fgraph = FunctionGraph([x, y, z], out)
    add_supervisor_to_fgraph(fgraph, [In(inp, mutable=True) for inp in fgraph.inputs])
    rewrite_graph(fgraph, include=("inplace",))
    assert fgraph.outputs[0].owner.op.destroy_map == {0: [0], 1: [1], 2: [2]}

    # Only x, z can be inplaced, y is protected
    out = blockwise_op(x.T, y.T, z.T)
    fgraph = FunctionGraph([x, y, z], out)
    add_supervisor_to_fgraph(
        fgraph, [In(inp, mutable=(i % 2) == 0) for i, inp in enumerate(fgraph.inputs)]
    )
    rewrite_graph(fgraph, include=("inplace",))
    assert fgraph.outputs[0].owner.op.destroy_map == {0: [0], 2: [2]}

    # Only y can be inplaced, x is reused for first and third outputs
    out = blockwise_op(x.T, y.T, x.T)
    fgraph = FunctionGraph([x, y, z], out)
    add_supervisor_to_fgraph(fgraph, [In(inp, mutable=True) for inp in fgraph.inputs])
    rewrite_graph(fgraph, include=("inplace",))
    assert fgraph.outputs[0].owner.op.destroy_map == {1: [1]}


def test_eig_blockwise():
    x = tensor("x", shape=(2, 3, 3), dtype="float64")
    eigen_values, eigen_vectors = eig(x)
    assert eigen_values.dtype == "complex128"
    assert eigen_vectors.dtype == "complex128"
    fn = function([x], [eigen_values, eigen_vectors])
    eigen_values_res, eigen_vectors_res = fn(np.full((2, 3, 3), np.eye(3)))
    np.testing.assert_allclose(
        eigen_values_res,
        np.ones((2, 3), dtype="complex128"),
        strict=True,
    )
    np.testing.assert_allclose(
        eigen_vectors_res,
        np.full((2, 3, 3), np.eye(3), dtype="complex128"),
        strict=True,
    )
