import numpy as np

from keras_core.src import backend
from keras_core.src.api_export import keras_core_export
from keras_core.src.layers.layer import Layer
from keras_core.src.utils import backend_utils
from keras_core.src.utils.module_utils import tensorflow as tf


@keras_core_export("keras_core.layers.RandomCrop")
class RandomCrop(Layer):
    """A preprocessing layer which randomly crops images during training.

    During training, this layer will randomly choose a location to crop images
    down to a target size. The layer will crop all the images in the same batch
    to the same cropping location.

    At inference time, and during training if an input image is smaller than the
    target size, the input will be resized and cropped so as to return the
    largest possible window in the image that matches the target aspect ratio.
    If you need to apply random cropping at inference time, set `training` to
    True when calling the layer.

    Input pixel values can be of any range (e.g. `[0., 1.)` or `[0, 255]`) and
    of integer or floating point dtype. By default, the layer will output
    floats.

    **Note:** This layer wraps `tf.keras.layers.RandomCrop`. It cannot
    be used as part of the compiled computation graph of a model with
    any backend other than TensorFlow.
    It can however be used with any backend when running eagerly.
    It can also always be used as part of an input preprocessing pipeline
    with any backend (outside the model itself), which is how we recommend
    to use this layer.

    **Note:** This layer is safe to use inside a `tf.data` pipeline
    (independently of which backend you're using).

    Input shape:
        3D (unbatched) or 4D (batched) tensor with shape:
        `(..., height, width, channels)`, in `"channels_last"` format.

    Output shape:
        3D (unbatched) or 4D (batched) tensor with shape:
        `(..., target_height, target_width, channels)`.

    Args:
        height: Integer, the height of the output shape.
        width: Integer, the width of the output shape.
        seed: Integer. Used to create a random seed.
        **kwargs: Base layer keyword arguments, such as
            `name` and `dtype`.
    """

    def __init__(self, height, width, seed=None, name=None, **kwargs):
        if not tf.available:
            raise ImportError(
                "Layer RandomCrop requires TensorFlow. "
                "Install it via `pip install tensorflow`."
            )

        super().__init__(name=name, **kwargs)
        self.seed = seed or backend.random.make_default_seed()
        self.layer = tf.keras.layers.RandomCrop(
            height=height,
            width=width,
            seed=self.seed,
            name=name,
        )
        self.supports_masking = False
        self.supports_jit = False
        self._convert_input_args = False
        self._allow_non_tensor_positional_args = True

    def call(self, inputs, training=True):
        if not isinstance(inputs, (tf.Tensor, np.ndarray, list, tuple)):
            inputs = tf.convert_to_tensor(backend.convert_to_numpy(inputs))
        outputs = self.layer.call(inputs, training=training)
        if (
            backend.backend() != "tensorflow"
            and not backend_utils.in_tf_graph()
        ):
            outputs = backend.convert_to_tensor(outputs)
        return outputs

    def compute_output_shape(self, input_shape):
        return tuple(self.layer.compute_output_shape(input_shape))

    def get_config(self):
        config = self.layer.get_config()
        config.update({"seed": self.seed})
        return config

