import keras
import numpy as np

from keras_hub.src.api_export import keras_hub_export
from keras_hub.src.models.backbone import Backbone


@keras_hub_export("keras_hub.models.ControlNetBackbone")
class ControlNetBackbone(Backbone):
    """Lightweight conditioning encoder for ControlNet.

    This backbone model takes a conditioning image (e.g., a Canny edge map) and
    encodes it into multi-scale feature maps. These feature maps can then be
    injected into the intermediate layers of a diffusion model like a UNet to
    condition its generation process.

    This implementation is a simplified version of the ControlNet encoder,
    consisting of a series of convolutional and pooling layers.

    Args:
        input_shape: tuple. The shape of the input image, excluding the batch
            dimension. Defaults to `(None, None, 1)`.
        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
            for model computations and weights.

    Example:
    ```python
    import numpy as np

    # Define the input for the conditioning image.
    input_data = {
        "pixel_values": np.ones(shape=(1, 512, 512, 1), dtype="float32"),
    }

    # Create a ControlNetBackbone instance.
    model = ControlNetBackbone()

    # Get the multi-scale feature maps.
    outputs = model(input_data)
    # `outputs` will be a dictionary with keys "scale_1", "scale_2", "scale_3".
    ```
    """

    def __init__(
        self,
        input_shape=(None, None, 1),
        dtype=None,
        **kwargs,
    ):
        # === Layers ===
        self.down1 = keras.layers.Conv2D(
            64, kernel_size=3, padding="same", activation="relu", name="down1_conv"
        )
        self.down2 = keras.layers.Conv2D(
            128, kernel_size=3, padding="same", activation="relu", name="down2_conv"
        )
        self.down3 = keras.layers.Conv2D(
            256, kernel_size=3, padding="same", activation="relu", name="down3_conv"
        )
        self.pool = keras.layers.MaxPooling2D(pool_size=2, name="pool")

        # === Functional Model ===
        pixel_values = keras.Input(
            shape=input_shape, dtype=dtype, name="pixel_values"
        )

        f1 = self.down1(pixel_values)
        p1 = self.pool(f1)

        f2 = self.down2(p1)
        p2 = self.pool(f2)

        f3 = self.down3(p2)

        outputs = {
            "scale_1": f1,
            "scale_2": f2,
            "scale_3": f3,
        }

        super().__init__(
            inputs=pixel_values,
            outputs=outputs,
            dtype=dtype,
            **kwargs,
        )

        # Store parameters as attributes.
        self.input_shape_arg = input_shape

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "input_shape": self.input_shape_arg,
            }
        )
        return config

import pytest
from keras import ops

from keras_hub.src.models.controlnet.controlnet_backbone import ControlNetBackbone
from keras_hub.src.tests.test_case import TestCase


class ControlNetBackboneTest(TestCase):
    def setUp(self):
        self.init_kwargs = {}
        self.input_data = {
            "pixel_values": ops.ones((2, 32, 32, 1), dtype="float32"),
        }

    def test_backbone_basics(self):
        self.run_backbone_test(
            cls=ControlNetBackbone,
            init_kwargs=self.init_kwargs,
            input_data=self.input_data,
            expected_output_shape={
                "scale_1": (2, 32, 32, 64),
                "scale_2": (2, 16, 16, 128),
                "scale_3": (2, 8, 8, 256),
            },
        )

    @pytest.mark.large
    def test_saved_model(self):
        self.run_model_saving_test(
            cls=ControlNetBackbone,
            init_kwargs=self.init_kwargs,
            input_data=self.input_data,
        )