from keras.src import backend
from keras.src import ops
from keras.src import testing
from keras.src.layers.core import input_layer
from keras.src.ops import operation_utils


class OperationUtilsTest(testing.TestCase):
    def test_get_source_inputs(self):
        x1 = backend.KerasTensor(shape=(2,))
        x2 = backend.KerasTensor(shape=(2,))
        x = x1 + x2
        x += 2
        x = ops.square(x)
        self.assertEqual(operation_utils.get_source_inputs(x), [x1, x2])

    def test_get_source_inputs_return_input_tensor(self):
        inputs = input_layer.Input(shape=(10,))
        self.assertIs(operation_utils.get_source_inputs(inputs)[0], inputs)

    def test_compute_expand_dims_output_shape(self):
        input_shape = (2, 3, 4)
        axis = -1
        output_shape = operation_utils.compute_expand_dims_output_shape(
            input_shape, axis
        )
        expected_output_shape = (2, 3, 4, 1)
        self.assertEqual(output_shape, expected_output_shape)

        input_shape = (2, 3, 4)
        axis = (1, -1)
        output_shape = operation_utils.compute_expand_dims_output_shape(
            input_shape, axis
        )
        expected_output_shape = (2, 1, 3, 4, 1)
        self.assertEqual(output_shape, expected_output_shape)

    def test_compute_pooling_output_shape(self):
        input_shape = (1, 4, 4, 1)
        pool_size = (2, 2)
        strides = (2, 2)
        output_shape = operation_utils.compute_pooling_output_shape(
            input_shape, pool_size, strides
        )
        expected_output_shape = (1, 2, 2, 1)
        self.assertEqual(output_shape, expected_output_shape)

    def test_compute_pooling_output_shape_with_none(self):
        input_shape = (None, 4, 4, 1)
        pool_size = (2, 2)
        strides = (2, 2)
        output_shape = operation_utils.compute_pooling_output_shape(
            input_shape, pool_size, strides
        )
        expected_output_shape = (None, 2, 2, 1)
        self.assertEqual(output_shape, expected_output_shape)

    def test_compute_pooling_output_shape_valid_padding(self):
        input_shape = (1, 4, 4, 1)
        pool_size = (2, 2)
        strides = (2, 2)
        output_shape = operation_utils.compute_pooling_output_shape(
            input_shape, pool_size, strides, padding="valid"
        )
        self.assertEqual(output_shape, (1, 2, 2, 1))

    def test_compute_pooling_output_shape_channels_last(self):
        input_shape = (1, 4, 4, 3)
        pool_size = (2, 2)
        strides = (2, 2)
        output_shape = operation_utils.compute_pooling_output_shape(
            input_shape,
            pool_size,
            strides,
            padding="valid",
            data_format="channels_last",
        )
        self.assertEqual(output_shape, (1, 2, 2, 3))

    def test_compute_pooling_output_shape_same_padding_stride1(self):
        input_shape = (1, 4, 4, 3)
        pool_size = (2, 2)
        strides = (1, 1)
        output_shape = operation_utils.compute_pooling_output_shape(
            input_shape,
            pool_size,
            strides,
            padding="same",
            data_format="channels_last",
        )
        self.assertEqual(output_shape, (1, 4, 4, 3))

    def test_compute_pooling_output_shape_rejects_zero_output(self):
        # pool_size > input spatial dim produces a zero-size output that the
        # downstream model cannot use. The Conv equivalent was fixed in #22418;
        # the pooling case should behave the same way.
        with self.assertRaisesRegex(
            ValueError, "Computed output size would be zero or negative"
        ):
            operation_utils.compute_pooling_output_shape(
                (1, 5, 5, 3), pool_size=(20, 20), strides=(20, 20)
            )

    def test_compute_conv_output_shape(self):
        input_shape = (1, 4, 4, 1)
        filters = 1
        kernel_size = (3, 3)
        strides = (1, 1)
        output_shape = operation_utils.compute_conv_output_shape(
            input_shape, filters, kernel_size, strides
        )
        expected_output_shape = (1, 2, 2, 1)
        self.assertEqual(output_shape, expected_output_shape)

    def test_compute_conv_output_shape_with_none(self):
        input_shape = (None, 4, 4, 1)
        kernel_size = (3, 3)
        filters = 1
        strides = (1, 1)
        output_shape = operation_utils.compute_conv_output_shape(
            input_shape, filters, kernel_size, strides
        )
        expected_output_shape = (None, 2, 2, 1)
        self.assertEqual(output_shape, expected_output_shape)

    def test_compute_conv_output_shape_valid_padding(self):
        input_shape = (1, 4, 4, 1)
        kernel_size = (3, 3)
        filters = 1
        strides = (2, 2)
        output_shape = operation_utils.compute_conv_output_shape(
            input_shape, filters, kernel_size, strides, padding="valid"
        )
        self.assertEqual(output_shape, (1, 1, 1, 1))

    def test_compute_conv_output_shape_channels_last(self):
        input_shape = (1, 4, 4, 3)
        kernel_size = (3, 3)
        filters = 3
        strides = (2, 2)
        output_shape = operation_utils.compute_conv_output_shape(
            input_shape,
            filters,
            kernel_size,
            strides,
            padding="valid",
            data_format="channels_last",
        )
        self.assertEqual(output_shape, (1, 1, 1, 3))

    def test_compute_conv_output_shape_same_padding_stride1(self):
        input_shape = (1, 4, 4, 3)
        kernel_size = (3, 3)
        filters = 3
        strides = (1, 1)
        output_shape = operation_utils.compute_conv_output_shape(
            input_shape,
            filters,
            kernel_size,
            strides,
            padding="same",
            data_format="channels_last",
        )
        self.assertEqual(output_shape, (1, 4, 4, 3))

    def test_validate_reshape_shape_rejects_invalid(self):
        with self.assertRaisesRegex(ValueError, "non-negative integer"):
            operation_utils.validate_reshape_shape((-2, 5))
        with self.assertRaisesRegex(
            ValueError, "at most one unknown dimension"
        ):
            operation_utils.validate_reshape_shape((-1, -1, 5))

    def test_validate_reshape_shape_accepts_dynamic_dims(self):
        # Non-int dimensions (e.g. backend tensor scalars such as
        # ``torch.SymInt`` produced by ``ops.shape(x)`` under
        # ``torch.compile``) must pass through without raising, even when
        # mixed with a single ``-1`` sentinel. This is the path
        # ``GroupNormalization._reshape_into_groups`` exercises and that
        # previously crashed dynamo tracing (see fix for #22868).
        class DynamicDim:
            """Stand-in for a non-int dynamic shape value."""

        d = DynamicDim()
        # Mixed int / -1 / dynamic — must not raise.
        operation_utils.validate_reshape_shape((-1, 32, d, d))
        operation_utils.validate_reshape_shape((d, d, d))
        operation_utils.validate_reshape_shape((d, -1, d, 8))
        # A second ``-1`` sentinel mixed with dynamic dims is still
        # rejected — dynamic dims do not absorb the ``-1`` count.
        with self.assertRaisesRegex(
            ValueError, "at most one unknown dimension"
        ):
            operation_utils.validate_reshape_shape((d, -1, d, -1, 4))

    def test_compute_reshape_output_shape(self):
        input_shape = (1, 4, 4, 1)
        target_shape = (16, 1)
        output_shape = operation_utils.compute_reshape_output_shape(
            input_shape, newshape=target_shape, newshape_arg_name="New shape"
        )
        self.assertEqual(output_shape, target_shape)

    def test_compute_reshape_output_shape_symbolic_tensors(self):
        # Entire newshape is a KerasTensor
        input_shape = (2, 10)
        shape_tensor = backend.KerasTensor((2,), dtype="int32")
        output_shape = operation_utils.compute_reshape_output_shape(
            input_shape, newshape=shape_tensor, newshape_arg_name="newshape"
        )
        self.assertEqual(output_shape, (None, None))

        # Tuple containing a KerasTensor and -1
        dim_tensor = backend.KerasTensor((), dtype="int32")
        output_shape2 = operation_utils.compute_reshape_output_shape(
            input_shape, newshape=(-1, dim_tensor), newshape_arg_name="newshape"
        )
        self.assertEqual(output_shape2, (None, None))

    def test_reduce_shape_no_axes_no_keepdims(self):
        input_shape = (1, 4, 4, 1)
        output_shape = operation_utils.reduce_shape(input_shape)
        expected_output_shape = ()
        self.assertEqual(output_shape, expected_output_shape)

    def test_reduce_shape_no_axes_with_keepdims(self):
        input_shape = (1, 4, 4, 1)
        output_shape = operation_utils.reduce_shape(input_shape, keepdims=True)
        expected_output_shape = (1, 1, 1, 1)
        self.assertEqual(output_shape, expected_output_shape)

    def test_reduce_shape_single_axis_no_keepdims(self):
        input_shape = (1, 4, 4, 1)
        axes = [1]
        output_shape = operation_utils.reduce_shape(input_shape, axes)
        expected_output_shape = (1, 4, 1)
        self.assertEqual(output_shape, expected_output_shape)

    def test_reduce_shape_single_axis_with_keepdims(self):
        input_shape = (1, 4, 4, 1)
        axes = [1]
        output_shape = operation_utils.reduce_shape(
            input_shape, axes, keepdims=True
        )
        expected_output_shape = (1, 1, 4, 1)
        self.assertEqual(output_shape, expected_output_shape)

    def test_reduce_shape_multiple_axes_no_keepdims(self):
        input_shape = (1, 4, 4, 1)
        axes = [1, 2]
        output_shape = operation_utils.reduce_shape(input_shape, axes)
        expected_output_shape = (1, 1)
        self.assertEqual(output_shape, expected_output_shape)

    def test_reduce_shape_out_of_order_axes_no_keepdims(self):
        input_shape = (1, 4, 4, 1)
        axes = [2, 1]
        output_shape = operation_utils.reduce_shape(input_shape, axes)
        expected_output_shape = (1, 1)
        self.assertEqual(output_shape, expected_output_shape)

    def test_reduce_shape_negative_axes_no_keepdims(self):
        input_shape = (1, 4, 4, 1)
        axes = [-2, -3]
        output_shape = operation_utils.reduce_shape(input_shape, axes)
        expected_output_shape = (1, 1)
        self.assertEqual(output_shape, expected_output_shape)
