import numpy as np

from keras.src import layers
from keras.src import testing


class AdaptivePooling2DLayerTest(testing.TestCase):
    """Tests for AdaptiveAveragePooling2D and AdaptiveMaxPooling2D."""

    def _run_layer_test(self, layer_class, x_np, output_size, data_format):
        """Helper: test layer output shape matches compute_output_shape()."""
        layer = layer_class(output_size=output_size, data_format=data_format)
        y = layer(x_np)
        expected_shape = layer.compute_output_shape(x_np.shape)
        self.assertEqual(y.shape, expected_shape)

    def test_average_pooling_basic_shapes(self):
        """Test AdaptiveAveragePooling2D basic shape transformation."""
        shape = (2, 3, 8, 8)  # N,C,H,W
        x = np.random.randn(*shape).astype("float32")
        self._run_layer_test(
            layers.AdaptiveAveragePooling2D,
            x,
            output_size=4,
            data_format="channels_first",
        )

    def test_max_pooling_basic_shapes(self):
        """Test AdaptiveMaxPooling2D basic shape transformation."""
        shape = (2, 3, 8, 8)
        x = np.random.randn(*shape).astype("float32")
        self._run_layer_test(
            layers.AdaptiveMaxPooling2D,
            x,
            output_size=4,
            data_format="channels_first",
        )

    def test_average_pooling_channels_last(self):
        """Test AdaptiveAveragePooling2D with channels_last format."""
        shape = (2, 8, 8, 3)  # N,H,W,C
        x = np.random.randn(*shape).astype("float32")
        self._run_layer_test(
            layers.AdaptiveAveragePooling2D,
            x,
            output_size=4,
            data_format="channels_last",
        )

    def test_max_pooling_channels_last(self):
        """Test AdaptiveMaxPooling2D with channels_last format."""
        shape = (2, 8, 8, 3)
        x = np.random.randn(*shape).astype("float32")
        self._run_layer_test(
            layers.AdaptiveMaxPooling2D,
            x,
            output_size=4,
            data_format="channels_last",
        )

    def test_average_pooling_tuple_output_size(self):
        """Test AdaptiveAveragePooling2D with tuple output_size."""
        shape = (2, 8, 8, 3)
        x = np.random.randn(*shape).astype("float32")
        self._run_layer_test(
            layers.AdaptiveAveragePooling2D,
            x,
            output_size=(4, 4),
            data_format="channels_last",
        )

    def test_max_pooling_tuple_output_size(self):
        """Test AdaptiveMaxPooling2D with tuple output_size."""
        shape = (2, 8, 8, 3)
        x = np.random.randn(*shape).astype("float32")
        self._run_layer_test(
            layers.AdaptiveMaxPooling2D,
            x,
            output_size=(2, 4),
            data_format="channels_last",
        )

    def test_average_pooling_compute_output_shape(self):
        """Test compute_output_shape() for AdaptiveAveragePooling2D."""
        layer = layers.AdaptiveAveragePooling2D(
            output_size=16, data_format="channels_last"
        )
        input_shape = (None, 64, 64, 3)
        output_shape = layer.compute_output_shape(input_shape)
        self.assertEqual(output_shape, (None, 16, 16, 3))

    def test_max_pooling_compute_output_shape(self):
        """Test compute_output_shape() for AdaptiveMaxPooling2D."""
        layer = layers.AdaptiveMaxPooling2D(
            output_size=(8, 16), data_format="channels_first"
        )
        input_shape = (2, 3, 64, 64)
        output_shape = layer.compute_output_shape(input_shape)
        self.assertEqual(output_shape, (2, 3, 8, 16))

    def test_average_pooling_get_config(self):
        """Test get_config() serialization for AdaptiveAveragePooling2D."""
        layer = layers.AdaptiveAveragePooling2D(
            output_size=32, data_format="channels_first"
        )
        config = layer.get_config()
        self.assertEqual(config["output_size"], (32, 32))
        self.assertEqual(config["data_format"], "channels_first")

    def test_max_pooling_get_config(self):
        """Test get_config() serialization for AdaptiveMaxPooling2D."""
        layer = layers.AdaptiveMaxPooling2D(
            output_size=(8, 16), data_format="channels_last"
        )
        config = layer.get_config()
        self.assertEqual(config["output_size"], (8, 16))
        self.assertEqual(config["data_format"], "channels_last")

    def test_average_pooling2d_numerical(self):
        """Test AdaptiveAveragePooling2D numerical correctness."""
        inputs = np.array(
            [
                [
                    [
                        [1.0, 2.0, 3.0, 4.0],
                        [5.0, 6.0, 7.0, 8.0],
                        [9.0, 10.0, 11.0, 12.0],
                        [13.0, 14.0, 15.0, 16.0],
                    ]
                ]
            ],
            dtype="float32",
        )
        expected = np.array([[[[3.5, 5.5], [11.5, 13.5]]]], dtype="float32")

        layer = layers.AdaptiveAveragePooling2D(
            output_size=2, data_format="channels_first"
        )
        outputs = layer(inputs)
        self.assertAllClose(outputs, expected, atol=1e-4)

    def test_max_pooling2d_numerical(self):
        """Test AdaptiveMaxPooling2D numerical correctness."""
        inputs = np.array(
            [
                [
                    [
                        [1.0, 2.0, 3.0, 4.0],
                        [5.0, 6.0, 7.0, 8.0],
                        [9.0, 10.0, 11.0, 12.0],
                        [13.0, 14.0, 15.0, 16.0],
                    ]
                ]
            ],
            dtype="float32",
        )
        expected = np.array([[[[6.0, 8.0], [14.0, 16.0]]]], dtype="float32")

        layer = layers.AdaptiveMaxPooling2D(
            output_size=2, data_format="channels_first"
        )
        outputs = layer(inputs)
        self.assertAllClose(outputs, expected, atol=1e-4)
