"""
TensorFlow functional helper utilities.
"""
from typing import Tuple

import tensorflow as tf  # type: ignore

from targetran._typing import ArrayLike
from targetran.utils import Interpolation

_TF_INTERPOLATION_DICT = {
    Interpolation.NEAREST: tf.image.ResizeMethod.NEAREST_NEIGHBOR,
    Interpolation.BILINEAR: tf.image.ResizeMethod.BILINEAR,
    Interpolation.BICUBIC: tf.image.ResizeMethod.BICUBIC,
    Interpolation.AREA: tf.image.ResizeMethod.AREA,
}


def _tf_convert(x: ArrayLike) -> tf.Tensor:
    if isinstance(x, tf.Tensor):
        return tf.cast(x, tf.float32)
    return tf.convert_to_tensor(x, tf.float32)


def _tf_cast_to_int(x: tf.Tensor) -> tf.Tensor:
    return tf.cast(x, tf.int32)


def _tf_round_to_int(x: tf.Tensor) -> tf.Tensor:
    return tf.cast(tf.math.rint(x), tf.int32)


def _tf_pad_image(
        image: tf.Tensor,
        pad_offsets: tf.Tensor,
) -> tf.Tensor:
    """
    pad_offsets: [top, bottom, left, right]
    """
    height = int(tf.shape(image)[0])
    width = int(tf.shape(image)[1])
    target_height = int(pad_offsets[0]) + height + int(pad_offsets[1])
    target_width = int(pad_offsets[2]) + width + int(pad_offsets[3])
    return tf.image.pad_to_bounding_box(
        image,
        int(pad_offsets[0]), int(pad_offsets[2]),
        target_height, target_width
    )


def _tf_resize_image(
        image: tf.Tensor,
        dest_size: Tuple[int, int],
        interpolation: Interpolation
) -> tf.Tensor:
    """
    dest_size: (image_height, image_width)
    """
    return tf.image.resize(
        image, size=dest_size, method=_TF_INTERPOLATION_DICT[interpolation]
    )


def _tf_gather_image(image: tf.Tensor, indices: tf.Tensor) -> tf.Tensor:
    """
    indices: [[row_0, row_1, ...], [col_0, col_1, ...]]
    """
    return tf.gather_nd(image, tf.transpose(indices))
