# coding=utf-8
# Copyright 2022 The Google Research Authors.
#
# 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.

"""Tests for tf3d.utils.box_utils."""

import numpy as np
import tensorflow as tf
from tf3d.utils import box_utils


class BoxUtilsTest(tf.test.TestCase):

  def test_point_features_to_box_features(self):
    # pyformat: disable
    points = tf.constant([[0.0, 0.0, 0.0],
                          [1.5, 0.0, 0.0],
                          [1.0, 1.5, 0.0],
                          [2.5, 2.5, 2.5],
                          [2.5, 2.5, 0.0]])
    point_features = tf.constant([[1.0, 0.0, 0.0, 0.0, 0.0],
                                  [0.0, 1.0, 0.0, 0.0, 0.0],
                                  [0.0, 0.0, 1.0, 0.0, 0.0],
                                  [0.0, 0.0, 0.0, 1.0, 0.0],
                                  [0.0, 0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)
    box_lengths = tf.constant([[4.0], [2.0]])
    box_heights = tf.constant([[1.0], [2.0]])
    box_widths = tf.constant([[2.0], [2.0]])
    box_rotations = tf.constant([[[1.0, 0.0, 0.0],
                                  [0.0, 1.0, 0.0],
                                  [0.0, 0.0, 1.0]],
                                 [[1.0, 0.0, 0.0],
                                  [0.0, 1.0, 0.0],
                                  [0.0, 0.0, 1.0]]])
    box_translations = tf.constant([[0.0, 0.0, 0.0],
                                    [1.0, 1.0, 1.0]])
    # pyformat: enable

    box_features = box_utils.point_features_to_box_features(
        points=points,
        point_features=point_features,
        num_bins_per_axis=2,
        boxes_length=box_lengths,
        boxes_height=box_heights,
        boxes_width=box_widths,
        boxes_rotation_matrix=box_rotations,
        boxes_center=box_translations)
    expected_box_features1 = tf.constant(
        [[0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0.5, 0.5, 0., 0., 0.]],
        dtype=tf.float32)
    expected_box_features2 = tf.constant(
        [[1., 0., 0., 0., 0.],
         [0., 1., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 1., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.]],
        dtype=tf.float32)
    expected_box_features = tf.reshape(
        tf.stack([expected_box_features1, expected_box_features2], axis=0),
        [2, -1])
    self.assertAllClose(box_features.numpy(), expected_box_features.numpy())

  def test_map_points_to_boxes_tf(self):
    # pyformat: disable
    points = tf.constant([[0.0, 0.0, 0.0],
                          [1.5, 0.0, 0.0],
                          [1.0, 1.5, 0.0],
                          [2.5, 2.5, 2.5],
                          [2.5, 2.5, 0.0]])
    box_lengths = tf.constant([[4.0], [1.0]])
    box_heights = tf.constant([[1.0], [1.0]])
    box_widths = tf.constant([[2.0], [1.0]])
    box_rotations = tf.constant([[[1.0, 0.0, 0.0],
                                  [0.0, 1.0, 0.0],
                                  [0.0, 0.0, 1.0]],
                                 [[1.0, 0.0, 0.0],
                                  [0.0, 1.0, 0.0],
                                  [0.0, 0.0, 1.0]]])
    box_translations = tf.constant([[0.0, 0.0, 0.0],
                                    [2.5, 2.5, 2.5]])
    # pyformat: enable
    box_indices = box_utils.map_points_to_boxes(
        points=points,
        boxes_length=box_lengths,
        boxes_height=box_heights,
        boxes_width=box_widths,
        boxes_rotation_matrix=box_rotations,
        boxes_center=box_translations,
        box_margin=0.0)
    self.assertAllEqual(box_indices.numpy(), np.array([0, 0, -1, 1, -1]))

  def test_get_box_corners_3d(self):
    # pyformat: disable
    box_lengths = tf.constant([[4.0], [1.0]])
    box_heights = tf.constant([[1.0], [1.0]])
    box_widths = tf.constant([[2.0], [1.0]])
    box_rotations = tf.constant([[[1.0, 0.0, 0.0],
                                  [0.0, 1.0, 0.0],
                                  [0.0, 0.0, 1.0]],
                                 [[1.0, 0.0, 0.0],
                                  [0.0, 1.0, 0.0],
                                  [0.0, 0.0, 1.0]]])
    box_translations = tf.constant([[0.0, 0.0, 0.0],
                                    [2.5, 2.5, 2.5]])
    # pyformat: enable
    box_corners = box_utils.get_box_corners_3d(
        boxes_length=box_lengths,
        boxes_height=box_heights,
        boxes_width=box_widths,
        boxes_rotation_matrix=box_rotations,
        boxes_center=box_translations)
    self.assertEqual(box_corners.shape, (2, 8, 3))
    expected_box_corners1 = np.array(
        [[2.0, 1.0, 0.5],
         [-2.0, 1.0, 0.5],
         [-2.0, -1.0, 0.5],
         [2.0, -1.0, 0.5],
         [2.0, 1.0, -0.5],
         [-2.0, 1.0, -0.5],
         [-2.0, -1.0, -0.5],
         [2.0, -1.0, -0.5]], np.float32)
    expected_box_corners2 = np.array(
        [[3.0, 3.0, 3.0],
         [2.0, 3.0, 3.0],
         [2.0, 2.0, 3.0],
         [3.0, 2.0, 3.0],
         [3.0, 3.0, 2.0],
         [2.0, 3.0, 2.0],
         [2.0, 2.0, 2.0],
         [3.0, 2.0, 2.0]], np.float32)
    self.assertAllClose(box_corners.numpy()[0], expected_box_corners1)
    self.assertAllClose(box_corners.numpy()[1], expected_box_corners2)

  def test_get_box_as_dotted_lines(self):
    corners = tf.repeat(