"""
Contains functions to compute ious of bounding boxes.

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

from keras_aug.datapoints import bounding_box


def _compute_area(box):
    """Computes area for bounding boxes

    Args:
      box: [N, 4] or [batch_size, N, 4] float Tensor, either batched
        or unbatched boxes.
    Returns:
      a float Tensor of [N] or [batch_size, N]
    """
    y_min, x_min, y_max, x_max = tf.split(
        box[..., :4], num_or_size_splits=4, axis=-1
    )
    return tf.squeeze((y_max - y_min) * (x_max - x_min), axis=-1)


def _compute_intersection(boxes1, boxes2):
    """Computes intersection area between two sets of boxes.

    Args:
      boxes1: [N, 4] or [batch_size, N, 4] float Tensor boxes.
      boxes2: [M, 4] or [batch_size, M, 4] float Tensor boxes.
    Returns:
      a [N, M] or [batch_size, N, M] float Tensor.
    """
    y_min1, x_min1, y_max1, x_max1 = tf.split(
        boxes1[..., :4], num_or_size_splits=4, axis=-1
    )
    y_min2, x_min2, y_max2, x_max2 = tf.split(
        boxes2[..., :4], num_or_size_splits=4, axis=-1
    )
    boxes2_rank = len(boxes2.shape)
    perm = [1, 0] if boxes2_rank == 2 else [0, 2, 1]
    # [N, M] or [batch_size, N, M]
    intersect_ymax = tf.minimum(y_max1, tf.transpose(y_max2, perm))
    intersect_ymin = tf.maximum(y_min1, tf.transpose(y_min2, perm))
    intersect_xmax = tf.minimum(x_max1, tf.transpose(x_max2, perm))
    intersect_xmin = tf.maximum(x_min1, tf.transpose(x_min2, perm))

    intersect_height = intersect_ymax - intersect_ymin
    intersect_width = intersect_xmax - intersect_xmin
    zeros_t = tf.cast(0, intersect_height.dtype)
    intersect_height = tf.maximum(zeros_t, intersect_height)
    intersect_width = tf.maximum(zeros_t, intersect_width)

    return intersect_height * intersect_width


def compute_iou(
    boxes1,
    boxes2,
    bounding_box_format,
    use_masking=False,
    mask_val=-1,
    images=None,
    image_shape=None,
):
    """Computes a lookup table vector containing the ious for a given set boxes.

    The lookup vector is to be indexed by [`boxes1_index`,`boxes2_index`] if
    boxes are unbatched and by [`batch`, `boxes1_index`,`boxes2_index`] if the
    boxes are batched.

    The users can pass `boxes1` and `boxes2` to be different ranks. For example:
    1) `boxes1`: [batch_size, M, 4], `boxes2`: [batch_size, N, 4] -> return
        [batch_size, M, N].
    2) `boxes1`: [batch_size, M, 4], `boxes2`: [N, 4] -> return
        [batch_size, M, N]
    3) `boxes1`: [M, 4], `boxes2`: [batch_size, N, 4] -> return
        [batch_size, M, N]
    4) `boxes1`: [M, 4], `boxes2`: [N, 4] -> return [M, N]

    Args:
      boxes1: a list of bounding boxes in 'corners' format. Can be batched or
        unbatched.
      boxes2: a list of bounding boxes in 'corners' format. Can be batched or
        unbatched.
      bounding_box_format: The format of bounding
            boxes of input dataset. Refer
            https://github.com/james77777778/keras-aug/blob/main/keras_aug/datapoints/bounding_box/converter.py
            for more details on supported bounding box formats.
      use_masking: whether masking will be applied. This will mask all `boxes1`
        or `boxes2` that have values less than 0 in all its 4 dimensions.
        Default to `False`.
      mask_val: int to mask those returned IOUs if the masking is True, defaults
        to -1.

    Returns:
      iou_lookup_table: a vector containing the pairwise ious of boxes1 and
        boxes2.
    """  # noqa: E501

    boxes1_rank = len(boxes1.shape)
    boxes2_rank = len(boxes2.shape)

    if boxes1_rank not in [2, 3]:
        raise ValueError(
            "compute_iou() expects boxes1 to be batched, or to be unbatched. "
            f"Received len(boxes1.shape)={boxes1_rank}, "
            f"len(boxes2.shape)={boxes2_rank}. Expected either "
            "len(boxes1.shape)=2 AND or len(boxes1.shape)=3."
        )
    if boxes2_rank not in [2, 3]:
        raise ValueError(
            "compute_iou() expects boxes2 to be batched, or to be unbatched. "
            f"Received len(boxes1.shape)={boxes1_rank}, "
            f"len(boxes2.shape)={boxes2_rank}. Expected either "
            "len(boxes2.shape)=2 AND or len(boxes2.shape)=3."
        )

    target_format = "yxyx"
    if bounding_box.is_relative(bounding_box_format):
        target_format = bounding_box.as_relative(target_format)

    boxes1 = bounding_box.convert_format(
        boxes1,
        source=bounding_box_format,
        target=target_format,
        images=images,
        image_shape=image_shape,
    )

    boxes2 = bounding_box.convert_format(
        boxes2,
        source=bounding_box_format,
        target=target_format,
        images=images,
        image_shape=image_shape,
    )

    intersect_area = _compute_intersection(boxes1, boxes2)
    boxes1_area = _compute_area(boxes1)
    boxes2_area = _compute_area(boxes2)
    boxes2_area_rank = len(boxes2_area.shape)
    boxes2_axis = 1 if (boxes2_area_rank == 2) else 0
    boxes1_area = tf.expand_dims(boxes1_area, axis=-1)
    boxes2_area = tf.expand_dims(boxes2_area, axis=boxes2_axis)
    union_area = boxes1_area + boxes2_area - intersect_area
    res = tf.math.divide_no_nan(intersect_area, union_area)

    if boxes1_rank == 2:
        perm = [1, 0]
    else:
        perm = [0, 2, 1]

    if not use_masking:
        return res

    mask_val_t = tf.cast(mask_val, res.dtype) * tf.ones_like(res)
    boxes1_mask = tf.less(tf.reduce_max(boxes1, axis=-1, keepdims=True), 0.0)
    boxes2_mask = tf.less(tf.reduce_max(boxes2, axis=-1, keepdims=True), 0.0)
    background_mask = tf.logical_or(
        boxes1_mask, tf.transpose(boxes2_mask, perm)
    )
    iou_lookup_table = tf.where(background_mask, mask_val_t, res)
    return iou_lookup_table
