"""
Most of these codes come from KerasCV.
"""
import tensorflow as tf
from keras import backend


def _unwrap_value_range(value_range, dtype=tf.float32):
    min_value, max_value = value_range
    min_value = tf.cast(min_value, dtype=dtype)
    max_value = tf.cast(max_value, dtype=dtype)
    return min_value, max_value


def transform_value_range(
    images, original_range, target_range, dtype=tf.float32
):
    """transforms values in input tensor from original_range to target_range.
    This function is intended to be used in preprocessing layers that
    rely upon color values. This allows us to assume internally that
    the input tensor is always in the range [0, 255].

    Args:
        images (tf.Tensor): the set of images to transform to the target range.
        original_range (Tuple[int|float, int|float]): The value range to
            transform from.
        target_range (Tuple[int|float, int|float]): The value range to
            transform to.
        dtype (tf.DType, optional): The dtype to compute the conversion with.
            Defaults to `tf.float32`.

    Returns:
        a new Tensor with values in the target range.

    References:
        - `KerasCV <https://github.com/keras-team/keras-cv>`_
    """
    if (
        original_range[0] == target_range[0]
        and original_range[1] == target_range[1]
    ):
        return images

    images = tf.cast(images, dtype=dtype)
    original_min_value, original_max_value = _unwrap_value_range(
        original_range, dtype=dtype
    )
    target_min_value, target_max_value = _unwrap_value_range(
        target_range, dtype=dtype
    )

    # images in the [0, 1] scale
    images = (images - original_min_value) / (
        original_max_value - original_min_value
    )

    scale_factor = target_max_value - target_min_value
    return (images * scale_factor) + target_min_value


def blend(images_1, images_2, factors, value_range=None):
    """Blend image1 and image2 using 'factors'. Can be batched inputs.

    Factor can be above ``0.0``.  A value of ``0.0`` means only image1 is used.
    A value of ``1.0`` means only image2 is used.  A value between ``0.0`` and
    ``1.0`` means we linearly interpolate the pixel values between the two
    images.  A value greater than ``1.0`` "extrapolates" the difference
    between the two pixel values. If ``value_range`` is set, the results will be
    clipped into ``value_range``

    Args:
        image1 (tf.Tensor): First image(s).
        image2 (tf.Tensor): Second image(s).
        factor (float|tf.Tensor): The blend factor(s).
        value_range (Sequence[int|float], optional): The value range of the
            results. Defaults to ``None``.

    References:
        - `KerasCV <https://github.com/keras-team/keras-cv>`_
    """
    results = images_1 + factors * (images_2 - images_1)
    if value_range is not None:
        results = tf.clip_by_value(results, value_range[0], value_range[1])
    return results


def rgb_to_grayscale(images):
    """Converts images from RGB to Grayscale.

    Compared to ``tf.image.rgb_to_grayscale``, this function replaces
    ``tf.tensordot`` with ``tf.math.multiply`` and ``tf.math.add`` to reduce
    memory usage.

    Args:
        images (tf.Tensor): The RGB tensor to convert. The last dimension must
            have size 3 and should contain RGB values.

    References:
        - `torchvision <https://github.com/pytorch/vision>`_
    """
    return (
        images[..., 0:1] * 0.2989
        + images[..., 1:2] * 0.587
        + images[..., 2:3] * 0.114
    )


def transform(
    images,
    transforms,
    fill_mode="reflect",
    fill_value=0.0,
    interpolation="bilinear",
    output_shape=None,
    name=None,
):
    """Applies the given transform(s) to the image(s).

    Args:
      images: A tensor of shape
        `(num_images, num_rows, num_columns, num_channels)` (NHWC). The rank
        must be statically known (the shape is not `TensorShape(None)`).
      transforms: Projective transform matrix/matrices. A vector of length 8 or
        tensor of size N x 8. If one row of transforms is
        [a0, a1, a2, b0, b1, b2, c0, c1], then it maps the *output* point
        `(x, y)` to a transformed *input* point
        `(x', y') = ((a0 x + a1 y + a2) / k, (b0 x + b1 y + b2) / k)`, where
        `k = c0 x + c1 y + 1`. The transforms are *inverted* compared to the
        transform mapping input points to output points. Note that gradients are
        not backpropagated into transformation parameters.
      fill_mode: Points outside the boundaries of the input are filled according
        to the given mode (one of `{"constant", "reflect", "wrap", "nearest"}`).
      fill_value: a float represents the value to be filled outside the
        boundaries when `fill_mode="constant"`.
      interpolation: Interpolation mode. Supported values: `"nearest"`,
        `"bilinear"`.
      output_shape: Output dimension after the transform, `[height, width]`.
        If `None`, output is the same size as input image.
      name: The name of the op.

    Fill mode behavior for each valid value is as follows:

    - reflect (d c b a | a b c d | d c b a)
    The input is extended by reflecting about the edge of the last pixel.

    - constant (k k k k | a b c d | k k k k)
    The input is extended by filling all
    values beyond the edge with the same constant value k = 0.

    - wrap (a b c d | a b c d | a b c d)
    The input is extended by wrapping around to the opposite edge.

    - nearest (a a a a | a b c d | d d d d)
    The input is extended by the nearest pixel.

    Input shape:
      4D tensor with shape: `(samples, height, width, channels)`,
        in `"channels_last"` format.

    Output shape:
      4D tensor with shape: `(samples, height, width, channels)`,
        in `"channels_last"` format.

    Returns:
      Image(s) with the same type and shape as `images`, with the given
      transform(s) applied. Transformed coordinates outside of the input image
      will be filled with zeros.

    Raises:
      TypeError: If `image` is an invalid type.
      ValueError: If output shape is not 1-D int32 Tensor.

    References:
        - `KerasCV <https://github.com/keras-team/keras-cv>`_
    """
    with backend.name_scope(name or "transform"):
        if output_shape is None:
            output_shape = tf.shape(images)[1:3]
            if not tf.executing_eagerly():
                output_shape_value = tf.get_static_value(output_shape)
                if output_shape_value is not None:
                    output_shape = output_shape_value

        output_shape = tf.convert_to_tensor(
            output_shape, tf.int32, name="output_shape"
        )

        if not output_shape.get_shape().is_compatible_with([2]):
            raise ValueError(
                "output_shape must be a 1-D Tensor of 2 elements: "
                "new_height, new_width, instead got "
                "{}".format(output_shape)
            )

        fill_value = tf.convert_to_tensor(
            fill_value, tf.float32, name="fill_value"
        )

        return tf.raw_ops.ImageProjectiveTransformV3(
            images=images,
            output_shape=output_shape,
            fill_value=fill_value,
            transforms=transforms,
            fill_mode=fill_mode.upper(),
            interpolation=interpolation.upper(),
        )
