# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Spatial transformation ops like RoIAlign, CropAndResize."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow.compat.v1 as tf


def _coordinate_vector_1d(start, end, size, align_endpoints):
  """Generates uniformly spaced coordinate vector.

  Args:
    start: A float tensor of shape [batch, num_boxes] indicating start values.
    end: A float tensor of shape [batch, num_boxes] indicating end values.
    size: Number of points in coordinate vector.
    align_endpoints: Whether to align first and last points exactly to
      endpoints.

  Returns:
    A 3D float tensor of shape [batch, num_boxes, size] containing grid
    coordinates.
  """
  start = tf.expand_dims(start, -1)
  end = tf.expand_dims(end, -1)
  length = end - start
  if align_endpoints:
    relative_grid_spacing = tf.linspace(0.0, 1.0, size)
    offset = 0 if size > 1 else length / 2
  else:
    relative_grid_spacing = tf.linspace(0.0, 1.0, size + 1)[:-1]
    offset = length / (2 * size)
  relative_grid_spacing = tf.reshape(relative_grid_spacing, [1, 1, size])
  relative_grid_spacing = tf.cast(relative_grid_spacing, dtype=start.dtype)
  absolute_grid = start + offset + relative_grid_spacing * length
  return absolute_grid


def box_grid_coordinate_vectors(boxes, size_y, size_x, align_corners=False):
  """Generates coordinate vectors for a `size x size` grid in boxes.

  Each box is subdivided uniformly into a grid consisting of size x s
ize
  rectangular cells. This function returns coordinate vectors describing
  the center of each cell.

  If `ali
gn_corners` is true, grid points are uniformly spread such that the
  corner points on the grid exactly overlap corners of the boxes.

  Note that output coordinates are expressed in the same coordinate frame as
  input boxes.

  Args:
    boxes: A float tensor of shape [batch, num_boxes, 4] containing boxes of the
      form [ymin, xmin, ymax, xmax].
    size_y: Size of the grid in y axis.
    size_x: Size of the grid in x axis.
    align_corners: Whether to align the corner grid points exactly with box
      corners.

  Returns:
    box_grid_y: A float tensor of shape [batch, num_boxes, size_y] containing y
      coordinates for grid points.
    box_grid_x: A float tensor of shape [batch, num_boxes, size_x] containing x
      coordinates for grid points.
  """
  ymin, xmin, ymax, xmax = tf.unstack(boxes, axis=-1)
  box_grid_y = _coordinate_vector_1d(ymin, ymax, size_y, align_corners)
  box_grid_x = _coordinate_vector_1d(xmin, xmax, size_x, align_corners)
  return box_grid_y, box_grid_x


def feature_grid_coordinate_vectors(box_grid_y, box_grid_x):
  """Returns feature grid point coordinate vectors for bilinear interpolation.

  Box grid is specified in absolute coordinate system with origin at left top
  (0, 0). The returned coordinate vectors contain 0-based feature point indices.

  This function snaps each point in the box grid to nearest 4 points on the
  feature map.

  In this function we also follow the convention of treating feature pixels as
  point objects with no spatial extent.

  Args:
    box_grid_y: A float tensor of shape [batch, num_boxes, size] containing y
      coordinate vector of the box grid.
    box_grid_x: A float tensor of shape [batch, num_boxes, size] containing x
      coordinate vector of the box grid.

  Returns:
    feature_grid_y0: An int32 tensor of shape [batch, num_boxes, size]
      containing y coordinate vector for the top neighbors.
    feature_grid_x0: A int32 tensor of shape [batch, num_boxes, size]
      containing x coordinate vector for the left neighbors.
    feature_grid_y1: A int32 tensor of shape [batch, num_boxes, size]
      containing y coordinate vector for the bottom neighbors.
    feature_grid_x1: A int32 tensor of shape [batch, num_boxes, size]
      containing x coordinate vector for the right neighbors.
  """
  feature_grid_y0 = tf.floor(box_grid_y)
  feature_grid_x0 = tf.floor(box_grid_x)
  feature_grid_y1 = tf.floor(box_grid_y + 1)
  feature_grid_x1 = tf.floor(box_grid_x + 1)
  feature_grid_y0 = tf.cast(feature_grid_y0, dtype=tf.int32)
  feature_grid_y1 = tf.cast(feature_grid_y1, dtype=tf.int32)
  feature_grid_x0 = tf.cast(feature_grid_x0, dtype=tf.int32)
  feature_grid_x1 = tf.cast(feature_grid_x1, dtype=tf.int32)
  return (feature_grid_y0, feature_grid_x0, feature_grid_y1, feature_grid_x1)


def _valid_indicator(feature_grid_y, feature_grid_x, true_feature_shapes):
  """Computes a indicator vector for valid indices.

  Computes an indicator vector which is true for points on feature map and
  false for points off feature map.

  Args:
    feature_grid_y: An int32 tensor of shape [batch, num_boxes, size_y]
      containing y coordinate vector.
    feature_grid_x: An int32 tensor of shape [batch, num_boxes, size_x]
      containing x coordinate vector.
    true_feature_shapes: A int32 tensor of shape [batch, num_boxes, 2]
      containing valid height and width of feature maps. Feature maps are
      assumed to be aligned to the left top corner.

  Returns:
    indices: A 1D bool tensor indicating valid feature indices.
  """
  height = tf.cast(true_feature_shapes[:, :, 0:1], dtype=feature_grid_y.dtype)
  width = tf.cast(true_feature_shapes[:, :, 1:2], dtype=feature_grid_x.dtype)
  valid_indicator = tf.logical_and(
      tf.expand_dims(
          tf.logical_and(feature_grid_y >= 0, tf.less(feature_grid_y, height)),
          3),
      tf.expand_dims(
          tf.logical_and(feature_grid_x >= 0, tf.less(feature_grid_x, width)),
          2))
  return tf.reshape(valid_indicator, [-1])


def ravel_indices(feature_grid_y, feature_grid_x, num_levels, height, width,
                  box_levels):
  """Returns grid indices in a flattened feature map of shape [-1, channels].

  The returned 1-D array can be used to gather feature grid points from a
  feature map that has been flattened from [batch, num_levels, max_height,
  max_width, channels] to [batch * num_levels * max_height * max_width,
  channels].

  Args:
    feature_grid_y: An int32 tensor of shape [batch, num_boxes, size_y]
      containing y coordinate vector.
    feature_grid_x: An int32 tensor of shape [batch, num_boxes, size_x]
      containing x coordinate vector.
    num_levels: Number of feature levels.
    height: An integer indicating the padded height of feature maps.
    width: An integer indicating the padded width of feature maps.
    box_levels: An int32 tensor of shape [batch, num_boxes] indicating
      feature level assigned to each box.

  Returns:
    indices: A 1D int32 tensor containing feature point indices in a flattened
      feature grid.
  """
  num_boxes = tf.shape(feature_grid_y)[1]
  batch_size = tf.shape(feature_grid_y)[0]
  size_y = tf.shape(feature_grid_y)[2]
  size_x = tf.shape(feature_grid_x)[2]
  height_dim_offset = width
  level_dim_offset = height * height_dim_offset
  batch_dim_offset = num_levels * level_dim_offset

  batch_dim_indices = (
      tf.reshape(
          tf.range(batch_size) * batch_dim_offset, [batch_size, 1, 1, 1]) *
      tf.ones([1, num_boxes, size_y, size_x], dtype=tf.int32))
  box_level_indices = (
      tf.reshape(box_levels * level_dim_offset, [batch_size, num_boxes, 1, 1]) *
      tf.ones([1, 1, size_y, size_x], dtype=tf.int32))
  height_indices = (
      tf.reshape(feature_grid_y * height_dim_offset,
                 [batch_size, num_boxes, size_y, 1]) *
      tf.ones([1, 1, 1, size_x],