import pytest
from absl.testing import parameterized

from keras.src import Sequential
from keras.src import backend
from keras.src import layers
from keras.src import ops
from keras.src import random
from keras.src import testing
from keras.src.backend.common.keras_tensor import KerasTensor


class ReshapeTest(testing.TestCase):
    @parameterized.named_parameters(
        [
            {"testcase_name": "dense", "sparse": False},
            {"testcase_name": "sparse", "sparse": True},
        ]
    )
    def test_reshape(self, sparse):
        if sparse and not backend.SUPPORTS_SPARSE_TENSORS:
            pytest.skip("Backend does not support sparse tensors.")

        self.run_layer_test(
            layers.Reshape,
            init_kwargs={"target_shape": (8, 1)},
            input_shape=(3, 2, 4),
            input_sparse=sparse,
            expected_output_shape=(3, 8, 1),
            expected_output_sparse=sparse,
            run_training_check=not sparse,
        )

        self.run_layer_test(
            layers.Reshape,
            init_kwargs={"target_shape": (8,)},
            input_shape=(3, 2, 4),
            input_sparse=sparse,
            expected_output_shape=(3, 8),
            expected_output_sparse=sparse,
            run_training_check=not sparse,
        )

        self.run_layer_test(
            layers.Reshape,
            init_kwargs={"target_shape": (2, 4)},
            input_shape=(3, 8),
            input_sparse=sparse,
            expected_output_shape=(3, 2, 4),
            expected_output_sparse=sparse,
            run_training_check=not sparse,
        )
        self.run_layer_test(
            layers.Reshape,
            init_kwargs={"target_shape": (-1, 1)},
            input_shape=(3, 2, 4),
            input_sparse=sparse,
            expected_output_shape=(3, 8, 1),
            expected_output_sparse=sparse,
            run_training_check=not sparse,
        )

        self.run_layer_test(
            layers.Reshape,
            init_kwargs={"target_shape": (1, -1)},
            input_shape=(3, 2, 4),
            input_sparse=sparse,
            expected_output_shape=(3, 1, 8),
            expected_output_sparse=sparse,
            run_training_check=not sparse,
        )

        self.run_layer_test(
            layers.Reshape,
            init_kwargs={"target_shape": (-1,)},
            input_shape=(3, 2, 4),
            input_sparse=sparse,
            expected_output_shape=(3, 8),
            expected_output_sparse=sparse,
            run_training_check=not sparse,
        )

        self.run_layer_test(
            layers.Reshape,
            init_kwargs={"target_shape": (2, -1)},
            input_shape=(3, 2, 4),
            input_sparse=sparse,
            expected_output_shape=(3, 2, 4),
            expected_output_sparse=sparse,
            run_training_check=not sparse,
        )

    def test_reshape_with_dynamic_batch_size(self):
        input_layer = layers.Input(shape=(2, 4))
        reshaped = layers.Reshape((8,))(input_layer)
        self.assertEqual(reshaped.shape, (None, 8))

    def test_reshape_rejects_invalid_target_shape(self):
        # Negative values other than -1 are invalid.
        with self.assertRaisesRegex(
            ValueError, "non-negative integer.*-1.*target_shape=\\(-2, 5\\)"
        ):
            layers.Reshape(target_shape=(-2, 5))

        # Existing constraint: at most one -1.
        with self.assertRaisesRegex(
            ValueError,
            "at most one unknown dimension.*target_shape=\\(-1, -1, 5\\)",
        ):
            layers.Reshape(target_shape=(-1, -1, 5))

    def test_reshape_with_dynamic_batch_size_and_minus_one(self):
        input = KerasTensor((None, 6, 4))
        layer = layers.Reshape((-1, 8))
        reshaped = backend.compute_output_spec(layer.__call__, input)
        self.assertEqual(reshaped.shape, (None, 3, 8))

    def test_reshape_layer_with_varying_input_size_and_minus_one(self):
        layer = layers.Reshape((-1, 8))
        res = layer(ops.ones((1, 6, 4), dtype="float32"))
        self.assertEqual(res.shape, (1, 3, 8))
        res = layer(ops.ones((1, 10, 4), dtype="float32"))
        self.assertEqual(res.shape, (1, 5, 8))

    def test_reshape_with_dynamic_dim_and_minus_one(self):
        input = KerasTensor((4, 6, None, 3))
        layer = layers.Reshape((-1, 3))
        reshaped = backend.compute_output_spec(layer.__call__, input)
        self.assertEqual(reshaped.shape, (4, None, 3))

    def test_reshape_sets_static_shape(self):
        input_layer = layers.Input(batch_shape=(2, None))
        reshaped = layers.Reshape((3, 5))(input_layer)
        # Also make sure the batch dim is not lost after reshape.
        self.assertEqual(reshaped.shape, (2, 3, 5))

    @pytest.mark.requires_trainable_backend
    def test_reshape_model_fit_with_varying_input_size_and_minus_one(self):
        def generator():
            yield (
                ops.ones((1, 12, 2), dtype="float32"),
                ops.zeros((1, 3, 8), dtype="float32"),
            )
            yield (
                ops.ones((1, 20, 2), dtype="float32"),
                ops.zeros((1, 5, 8), dtype="float32"),
            )

        layer = layers.Reshape((-1, 8))
        model = Sequential([layer])
        model.compile(loss="mean_squared_error")
        model.fit(generator(), steps_per_epoch=2, epochs=1)

        # Also test inference with varying shapes to ensure -1 works dynamically
        # Input: (batch, seq_len, 2), Output: (batch, seq_len*2/8, 8)
        for seq_len, expected_first_dim in [(12, 3), (20, 5), (24, 6)]:
            x_test = random.normal((1, seq_len, 2))
            output = model(x_test)
            # Verify output shape is correct
            self.assertEqual(output.shape[1], expected_first_dim)
            self.assertEqual(output.shape[2], 8)
