import contextlib

import numpy as np
import pytest

import pytensor
import pytensor.tensor as pt
import pytensor.tensor.math as ptm
from pytensor import config, function
from pytensor.compile import get_mode
from pytensor.compile.ops import deep_copy_op
from pytensor.gradient import grad
from pytensor.scalar import Composite, float64
from pytensor.scalar import add as scalar_add
from pytensor.tensor import blas, matrix, tensor3
from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
from pytensor.tensor.math import All, Any, Max, Min, Prod, ProdWithoutZeros, Sum
from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad
from tests.link.numba.test_basic import (
    compare_numba_and_py,
    numba_mode,
    scalar_my_multi_out,
)
from tests.tensor.test_elemwise import check_elemwise_runtime_broadcast


rng = np.random.default_rng(42849)

add_inplace = Elemwise(scalar_add, {0: 0})


@pytest.mark.parametrize(
    "inputs, input_vals, output_fn",
    [
        (
            [pt.vector()],
            [rng.uniform(size=100).astype(config.floatX)],
            lambda x: pt.gammaln(x),
        ),
        (
            [pt.vector()],
            [rng.standard_normal(100).astype(config.floatX)],
            lambda x: pt.sigmoid(x),
        ),
        (
            [pt.vector()],
            [rng.standard_normal(100).astype(config.floatX)],
            lambda x: pt.log1mexp(x),
        ),
        (
            [pt.vector()],
            [rng.standard_normal(100).astype(config.floatX)],
            lambda x: pt.erf(x),
        ),
        (
            [pt.vector()],
            [rng.standard_normal(100).astype(config.floatX)],
            lambda x: pt.erfc(x),
        ),
        (
            [pt.vector()],
            [rng.standard_normal(100).astype(config.floatX)],
            lambda x: pt.erfcx(x),
        ),
        (
            [pt.vector() for i in range(4)],
            [rng.standard_normal(100).astype(config.floatX) for i in range(4)],
            lambda x, y, x1, y1: (x + y) * (x1 + y1) * y,
        ),
        (
            [pt.matrix(), pt.scalar()],
            [rng.normal(size=(2, 2)).astype(config.floatX), 0.0],
            lambda a, b: pt.switch(a, b, a),
        ),
        (
            [pt.scalar(), pt.scalar()],
            [
                np.array(1.0, dtype=config.floatX),
                np.array(1.0, dtype=config.floatX),
            ],
            lambda x, y: add_inplace(deep_copy_op(x), deep_copy_op(y)),
        ),
        (
            [pt.vector(), pt.vector()],
            [
                rng.standard_normal(100).astype(config.floatX),
                rng.standard_normal(100).astype(config.floatX),
            ],
            lambda x, y: add_inplace(deep_copy_op(x), deep_copy_op(y)),
        ),
        (
            [pt.vector(), pt.vector()],
            [
                rng.standard_normal(100).astype(config.floatX),
                rng.standard_normal(100).astype(config.floatX),
            ],
            lambda x, y: scalar_my_multi_out(x, y),
        ),
    ],
    ids=[
        "gammaln",
        "sigmoid",
        "log1mexp",
        "erf",
        "erfc",
        "erfcx",
        "complex_arithmetic",
        "switch",
        "add_inplace_scalar",
        "add_inplace_vector",
        "scalar_multi_out",
    ],
)
def test_Elemwise(inputs, input_vals, output_fn):
    outputs = output_fn(*inputs)
    if not isinstance(outputs, tuple | list):
        outputs = [outputs]

    compare_numba_and_py(
        inputs,
        outputs,
        input_vals,
        inplace=outputs[0].owner.op.destroy_map,
    )


@pytest.mark.xfail(reason="Logic had to be reversed due to surprising segfaults")
def test_elemwise_runtime_broadcast():
    check_elemwise_runtime_broadcast(get_mode("NUMBA"))


@pytest.mark.parametrize(
    "v, new_order",
    [
        # `{'drop': [], 'shuffle': [], 'augment': [0, 1]}`
        (
            (
                pt.lscalar(name="a"),
                np.array(1, dtype=np.int64),
            ),
            ("x", "x"),
        ),
        # I.e. `a_pt.T`
        # `{'drop': [], 'shuffle': [1, 0], 'augment': []}`
        (
            (pt.matrix("a"), np.array([[1.0, 2.0], [3.0, 4.0]], dtype=config.floatX)),
            (1, 0),
        ),
        # `{'drop': [], 'shuffle': [0, 1], 'augment': [2]}`
        (
            (pt.matrix("a"), np.array([[1.0, 2.0], [3.0, 4.0]], dtype=config.floatX)),
            (1, 0, "x"),
        ),
        # `{'drop': [1], 'shuffle': [2, 0], 'augment': [0, 2, 4]}`
        (
            (
                pt.tensor(dtype=config.floatX, shape=(None, 1, None), name="a"),
                np.array([[[1.0, 2.0]], [[3.0, 4.0]]], dtype=config.floatX),
            ),
            ("x", 2, "x", 0, "x"),
        ),
        # I.e. `a_pt.dimshuffle((0,))`
        # `{'drop': [1], 'shuffle': [0], 'augment': []}`
        (
            (
                pt.tensor(dtype=config.floatX, shape=(None, 1), name="a"),
                np.array([[1.0], [2.0], [3.0], [4.0]], dtype=config.floatX),
            ),
            (0,),
        ),
        (
            (
                pt.tensor(dtype=config.floatX, shape=(None, 1), name="a"),
                np.array([[1.0], [2.0], [3.0], [4.0]], dtype=config.floatX),
            ),
            (0,),
        ),
        (
            (
                pt.tensor(dtype=config.floatX, shape=(1, 1, 1), name="a"),
                np.array([[[1.0]]], dtype=config.floatX),
            ),
            (),
        ),
    ],
)
def test_Dimshuffle(v, new_order):
    v, v_test_value = v
    g = v.dimshuffle(new_order)
    compare_numba_and_py(
        [v],
        [g],
        [v_test_value],
    )


def test_Dimshuffle_returns_array():
    x = pt.vector("x", shape=(1,))
    y = 2 * x.dimshuffle([])
    func = pytensor.function([x], y, mode="NUMBA")
    out = func(np.zeros(1, dtype=config.floatX))
    assert out.ndim == 0


def test_Dimshuffle_non_contiguous():
    """The numba impl of reshape doesn't work with
    non-contiguous arrays, make sure we work around that."""
    x = pt.dvector()
    idx = pt.vector(dtype="int64")
    op = DimShuffle(input_ndim=1, new_order=[])
    out = op(pt.specify_shape(x[idx][::2], (1,)))
    func = pytensor.function([x, idx], out, mode="NUMBA")
    assert func(np.zeros(3), np.array([1])).ndim == 0


def test_Dimshuffle_squeeze_errors():
    x = pt.tensor3("x", shape=(4, None, 5))
    out = pt.squeeze(x, axis=1)
    assert out.type.shape == (4, 5)
    fn = function([x], out, mode=numba_mode)
    with pytest.raises(
        ValueError, match="Attempting to squeeze axes with size not equal to one"
    ):
        fn(np.zeros((4, 2, 5)))


@pytest.mark.parametrize(
    "careduce_fn, axis, v",
    [
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Sum(
                axis=axis, dtype=dtype, acc_dtype=acc_dtype
            )(x),
            0,
            (pt.vector(), np.arange(3, dtype=config.floatX)),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: All(axis)(x),
            0,
            (pt.vector(dtype="bool"), np.array([False, True, False])),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Any(axis)(x),
            0,
            (pt.vector(dtype="bool"), np.array([False, True, False])),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Sum(
                axis=axis, dtype=dtype, acc_dtype=acc_dtype
            )(x),
            0,
            (pt.matrix(), np.arange(3 * 2, dtype=config.floatX).reshape((3, 2))),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Sum(
                axis=axis, dtype=dtype, acc_dtype=acc_dtype
            )(x),
            (0, 1),
            (pt.matrix(), np.arange(3 * 2, dtype=config.floatX).reshape((3, 2))),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Sum(
                axis=axis, dtype=dtype, acc_dtype=acc_dtype
            )(x),
            (1, 0),
            (pt.matrix(), np.arange(3 * 2, dtype=config.floatX).reshape((3, 2))),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Sum(
                axis=axis, dtype=dtype, acc_dtype=acc_dtype
            )(x),
            None,
            (pt.matrix(), np.arange(3 * 2, dtype=config.floatX).reshape((3, 2))),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Sum(
                axis=axis, dtype=dtype, acc_dtype=acc_dtype
            )(x),
            1,
            (pt.matrix(), np.arange(3 * 2, dtype=config.floatX).reshape((3, 2))),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Prod(
                axis=axis, dtype=dtype, acc_dtype=acc_dtype
            )(x),
            (),  # Empty axes would normally be rewritten away, but we want to test it still works
            (pt.matrix(), np.arange(3 * 2, dtype=config.floatX).reshape((3, 2))),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Prod(
                axis=axis, dtype=dtype, acc_dtype=acc_dtype
            )(x),
            None,
            (
                pt.scalar(),
                np.array(99.0, dtype=config.floatX),
            ),  # Scalar input would normally be rewritten away, but we want to test it still works
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Prod(
                axis=axis, dtype=dtype, acc_dtype=acc_dtype
            )(x),
            0,
            (pt.vector(), np.arange(3, dtype=config.floatX)),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: ProdWithoutZeros(
                axis=axis, dtype=dtype, acc_dtype=acc_dtype
            )(x),
            0,
            (pt.vector(), np.arange(3, dtype=config.floatX)),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Prod(
                axis=axis, dtype=dtype, acc_dtype=acc_dtype
            )(x),
            0,
            (pt.matrix(), np.arange(3 * 2, dtype=config.floatX).reshape((3, 2))),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Prod(
                axis=axis, dtype=dtype, acc_dtype=acc_dtype
            )(x),
            1,
            (pt.matrix(), np.arange(3 * 2, dtype=config.floatX).reshape((3, 2))),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Max(axis)(x),
            None,
            (pt.matrix(), np.arange(3 * 2, dtype=config.floatX).reshape((3, 2))),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Max(axis)(x),
            None,
            (pt.lmatrix(), np.arange(3 * 2, dtype=np.int64).reshape((3, 2))),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Min(axis)(x),
            None,
            (pt.matrix(), np.arange(3 * 2, dtype=config.floatX).reshape((3, 2))),
        ),
        (
            lambda x, axis=None, dtype=None, acc_dtype=None: Min(axis)(x),
            None,
            (pt.lmatrix(), np.arange(3 * 2, dtype=np.int64).reshape((3, 2))),
        ),
    ],
)
def test_CAReduce(careduce_fn, axis, v):
    v, v_test_value = v
    g = careduce_fn(v, axis=axis)

    fn, _ = compare_numba_and_py(
        [v],
        [g],
        [v_test_value],
    )
    # Confirm CAReduce is in the compiled function
    # fn.dprint()
    [node] = fn.maker.fgraph.apply_nodes
    assert isinstance(node.op, CAReduce)


@pytest.mark.parametrize("axis", (-1, (0, -1), None))
def test_CAReduce_respects_acc_dtype(axis):
    x = tensor3("x", dtype="int8")
    out = x.sum(dtype="int8", acc_dtype="int64", axis=axis)
    # Choose values that would overflow if accumulated internally in int8
    max_int8 = np.iinfo(np.int8).max
    test_x = np.array([max_int8, 5, max_int8, -max_int8, 5, -max_int8], dtype=np.int8)
    test_x = np.broadcast_to(test_x, (6, 2, 6)).copy()
    _, [res] = compare_numba_and_py(
        [x],
        [out],
        [test_x],
    )
    if axis == -1:
        assert np.all(res == 10)
    elif axis == (0, -1):
        assert np.all(res == 60)
    elif axis is None:
        assert res == 120


@pytest.mark.parametrize("axis", (1, None))
def test_CAReduce_acc_complex_out_float(axis):
    x = matrix("x", dtype="complex128")
    out = x.sum(dtype="float64", axis=axis)
    test_x = np.array([[1 + 0.5j, 2 - 0.5j], [3 + 0.5j, 4 - 0.5j]], dtype="complex128")
    compare_numba_and_py([x], [out], [test_x])


@pytest.mark.parametrize("axis", (-1, (0, -1), None))
def test_CAReduce_discrete_infinity_identity(axis):
    rng = np.random.default_rng(337)
    x = tensor3("x", dtype="int8")
    out = x.max(axis)
    compare_numba_and_py(
        [x], [out], [rng.integers(-127, 127, size=(6, 6, 6)).astype("int8")]
    )


def test_scalar_Elemwise_Clip():
    a = pt.scalar("a")
    b = pt.scalar("b")
    inputs = [a, b]

    z = pt.switch(1, a, b)
    c = pt.clip(z, 1, 3)

    compare_numba_and_py(inputs, [c], [1, 1])


@pytest.mark.parametrize(
    "dy, sm, axis, exc",
    [
        (
            (pt.matrix(), np.array([[1, 1, 1], [0, 0, 0]], dtype=config.floatX)),
            (pt.matrix(), rng.random(size=(2, 3)).astype(config.floatX)),
            None,
            None,
        ),
        (
            (pt.matrix(), np.array([[1, 1, 1], [0, 0, 0]], dtype=config.floatX)),
            (pt.matrix(), rng.random(size=(2, 3)).astype(config.floatX)),
            0,
            None,
        ),
        (
            (pt.matrix(), np.array([[1, 1, 1], [0, 0, 0]], dtype=config.floatX)),
            (pt.matrix(), rng.random(size=(2, 3)).astype(config.floatX)),
            1,
            None,
        ),
    ],
)
def test_SoftmaxGrad(dy, sm, axis, exc):
    dy, dy_test_value = dy
    sm, sm_test_value = sm
    g = SoftmaxGrad(axis=axis)(dy, sm)

    cm = contextlib.suppress() if exc is None else pytest.warns(exc)
    with cm:
        compare_numba_and_py(
            [dy, sm],
            [g],
            [dy_test_value, sm_test_value],
        )


def test_SoftMaxGrad_constant_dy():
    dy = pt.constant(np.zeros((3,), dtype=config.floatX))
    sm = pt.vector(shape=(3,))
    inputs = [sm]

    g = SoftmaxGrad(axis=None)(dy, sm)

    compare_numba_and_py(inputs, [g], [np.ones((3,), dtype=config.floatX)])


@pytest.mark.parametrize(
    "x, axis, exc",
    [
        (
            (pt.vector(), rng.random(size=(2,)).astype(config.floatX)),
            None,
            None,
        ),
        (
            (pt.matrix(), rng.random(size=(2, 3)).astype(config.floatX)),
            None,
            None,
        ),
        (
            (pt.matrix(), rng.random(size=(2, 3)).astype(config.floatX)),
            0,
            None,
        ),
    ],
)
def test_Softmax(x, axis, exc):
    x, x_test_value = x
    g = Softmax(axis=axis)(x)

    cm = contextlib.suppress() if exc is None else pytest.warns(exc)
    with cm:
        compare_numba_and_py(
            [x],
            [g],
            [x_test_value],
        )


@pytest.mark.parametrize(
    "x, axis, exc",
    [
        (
            (pt.vector(), rng.random(size=(2,)).astype(config.floatX)),
            None,
            None,
        ),
        (
            (pt.matrix(), rng.random(size=(2, 3)).astype(config.floatX)),
            0,
            None,
        ),
        (
            (pt.matrix(), rng.random(size=(2, 3)).astype(config.floatX)),
            1,
            None,
        ),
    ],
)
def test_LogSoftmax(x, axis, exc):
    x, x_test_value = x
    g = LogSoftmax(axis=axis)(x)

    cm = contextlib.suppress() if exc is None else pytest.warns(exc)
    with cm:
        compare_numba_and_py(
            [x],
            [g],
            [x_test_value],
        )


@pytest.mark.parametrize(
    "x, axes, exc",
    [
        (
            (pt.dscalar(), np.array(0.0, dtype="float64")),
            [],
            None,
        ),
        (
            (pt.dvector(), rng.random(size=(3,)).astype("float64")),
            [0],
            None,
        ),
        (
            (pt.dmatrix(), rng.random(size=(3, 2)).astype("float64")),
            [0],
            None,
        ),
        (
            (pt.dmatrix(), rng.random(size=(3, 2)).astype("float64")),
            [0, 1],
            None,
        ),
    ],
)
def test_Max(x, axes, exc):
    x, x_test_value = x
    g = ptm.Max(axes)(x)

    cm = contextlib.suppress() if exc is None else pytest.warns(exc)
    with cm:
        compare_numba_and_py(
            [x],
            [g],
            [x_test_value],
        )


@pytest.mark.parametrize(
    "x, axes, exc",
    [
        (
            (pt.dscalar(), np.array(0.0, dtype="float64")),
            [],
            None,
        ),
        (
            (pt.dvector(), rng.random(size=(3,)).astype("float64")),
            [0],
            None,
        ),
        (
            (pt.dmatrix(), rng.random(size=(3, 2)).astype("float64")),
            [0],
            None,
        ),
        (
            (pt.dmatrix(), rng.random(size=(3, 2)).astype("float64")),
            [0, 1],
            None,
        ),
        (
            (pt.dmatrix(), rng.random(size=(3, 2)).astype("float64")),
            None,
            None,
        ),
    ],
)
def test_Argmax(x, axes, exc):
    x, x_test_value = x
    g = ptm.Argmax(axes)(x)

    cm = contextlib.suppress() if exc is None else pytest.warns(exc)
    with cm:
        compare_numba_and_py(
            [x],
            [g],
            [x_test_value],
        )


def test_elemwise_inplace_out_type():
    # Create a graph with an elemwise
    # Ravel failes if the elemwise output type is reported incorrectly
    x = pt.matrix()
    y = (2 * x).ravel()

    # Pass in the input as mutable, to trigger the inplace rewrites
    func = pytensor.function([pytensor.In(x, mutable=True)], y, mode="NUMBA")

    # Apply it to a numpy array that is neither C or F contigous
    x_val = np.broadcast_to(np.zeros((3,)), (6, 3))

    assert func(x_val).shape == (18,)


def test_elemwise_multiple_inplace_outs():
    x = pt.vector()
    y = pt.vector()

    x_ = pt.scalar_from_tensor(x[0])
    y_ = pt.scalar_from_tensor(y[0])
    out_ = x_ + 1, y_ + 1

    composite_op = Composite([x_, y_], out_)
    elemwise_op = Elemwise(composite_op, inplace_pattern={0: 0, 1: 1})
    out = elemwise_op(x, y)

    fn = function([x, y], out, mode="NUMBA", accept_inplace=True)
    x_test = np.array([1, 2, 3], dtype=config.floatX)
    y_test = np.array([4, 5, 6], dtype=config.floatX)
    out1, out2 = fn(x_test, y_test)
    assert out1 is x_test
    assert out2 is y_test
    np.testing.assert_allclose(out1, [2, 3, 4])
    np.testing.assert_allclose(out2, [5, 6, 7])


def test_scalar_loop():
    a_scalar = float64("a")
    const_scalar = float64("const")
    scalar_loop = pytensor.scalar.ScalarLoop(
        init=[a_scalar],
        update=[a_scalar + a_scalar + const_scalar],
        constant=[const_scalar],
    )

    a = pt.tensor("a", shape=(3,))
    const = pt.tensor("const", shape=(3,))
    n_steps = 3
    elemwise_loop = Elemwise(scalar_loop)(n_steps, a, const)

    compare_numba_and_py(
        [a, const],
        [elemwise_loop],
        [np.array([1, 2, 3], dtype="float64"), np.array([1, 1, 1], dtype="float64")],
    )


def test_gammainc_wrt_k_grad():
    x = pt.vector("x", dtype="float64")
    k = pt.vector("k", dtype="float64")

    out = pt.gammainc(k, x)
    grad_out = grad(out.sum(), k)

    compare_numba_and_py(
        [x, k],
        [grad_out],
        # These values of x and k trigger all the branches in the gradient of gammainc
        [
            np.array([0.0, 29.0, 31.0], dtype="float64"),
            np.array([1.0, 13.0, 11.0], dtype="float64"),
        ],
    )


@pytest.mark.parametrize(
    "x, y",
    [
        (
            (pt.matrix(), rng.random(size=(3, 2)).astype(config.floatX)),
            (pt.vector(), rng.random(size=(2,)).astype(config.floatX)),
        ),
        (
            (pt.matrix(dtype="float64"), rng.random(size=(3, 2)).astype("float64")),
            (pt.vector(dtype="float32"), rng.random(size=(2,)).astype("float32")),
        ),
        (
            (pt.lmatrix(), rng.poisson(size=(3, 2))),
            (pt.fvector(), rng.random(size=(2,)).astype("float32")),
        ),
        (
            (pt.lvector(), rng.random(size=(2,)).astype(np.int64)),
            (pt.lvector(), rng.random(size=(2,)).astype(np.int64)),
        ),
        (
            (pt.vector(dtype="int16"), rng.random(size=(2,)).astype(np.int16)),
            (pt.vector(dtype="uint8"), rng.random(size=(2,)).astype(np.uint8)),
        ),
        # Viewing the array with 2 last dimensions as complex128 means
        # the first entry will be real part and the second entry the imaginary part
        (
            (
                pt.matrix(dtype="complex128"),
                rng.random(size=(5, 4, 2)).view("complex128").squeeze(-1),
            ),
            (
                pt.matrix(dtype="complex128"),
                rng.random(size=(4, 3, 2)).view("complex128").squeeze(-1),
            ),
        ),
        (
            (pt.matrix(dtype="int64"), rng.random(size=(5, 4)).astype("int64")),
            (
                pt.matrix(dtype="complex128"),
                rng.random(size=(4, 3, 2)).view("complex128").squeeze(-1),
            ),
        ),
    ],
)
def test_Dot(x, y):
    x, x_test_value = x
    y, y_test_value = y

    g = ptm.dot(x, y)

    compare_numba_and_py(
        [x, y],
        [g],
        [x_test_value, y_test_value],
    )


@pytest.mark.parametrize(
    "x, y, exc",
    [
        (
            (
                pt.dtensor3(),
                rng.random(size=(2, 3, 3)).astype("float64"),
            ),
            (
                pt.dtensor3(),
                rng.random(size=(2, 3, 3)).astype("float64"),
            ),
            None,
        ),
        (
            (
                pt.dtensor3(),
                rng.random(size=(2, 3, 3)).astype("float64"),
            ),
            (
                pt.ltensor3(),
                rng.poisson(size=(2, 3, 3)).astype("int64"),
            ),
            None,
        ),
    ],
)
def test_BatchedDot(x, y, exc):
    x, x_test_value = x
    y, y_test_value = y

    g = blas.BatchedDot()(x, y)

    cm = contextlib.suppress() if exc is None else pytest.warns(exc)
    with cm:
        compare_numba_and_py(
            [x, y],
            g,
            [x_test_value, y_test_value],
        )
