# Copyright 2022 The sunds 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 isun.geometry.cameras."""

import numpy as np
from sunds.core.tf_geometry import cameras
import tensorflow as tf


class PinholeCameraTest(tf.test.TestCase):

  def test_construct_from_fov(self):
    """Construct camera using field of view and verify camera attributes."""
    camera_model = cameras.PinholeCamera.from_fov(
        image_width=640,
        image_height=480,
        horizontal_fov_in_degrees=57.5,
        vertical_fov_in_degrees=45.0)
    self.assertEqual(camera_model.image_width, 640)
    self.assertEqual(camera_model.image_height, 480)
    self.assertAllClose(
        camera_model.K,
        [[583.28296, 0., 320.], [0., 579.41125, 240.], [0., 0., 1.]])

  def test_construct_from_intrinsics(self):
    """Construct camera using intrinsics params and verify camera attributes."""
    camera_model = cameras.PinholeCamera.from_intrinsics(
        image_size_in_pixels=(640, 480),
        focal_length_in_pixels=(500.0, 400.0),
        principal_point_in_pixels=(320.1, 239.9),
        skew=0.1)
    self.assertEqual(camera_model.image_width, 640)
    self.assertEqual(camera_model.image_height, 480)
    self.assertAllClose(camera_model.K,
                        [[500.0, 0.1, 320.1], [0., 400.0, 239.9], [0., 0., 1.]])

  def test_project_principal_point_single(self):
    """3D points on +Z axis should project to principal point of the camera."""
    image_width, image_height = 640, 480
    fx, fy, cx, cy = 500.0, 500.0, image_width / 2, image_height / 2
    camera_model = cameras.PinholeCamera(
        K=np.array([[fx, 0., cx], [0., fy, cy], [0., 0., 1.]]),
        image_width=image_width,
        image_height=image_height)

    points_camera = np.array([0., 0., 1.])
    points_image = camera_model.project(points_camera)
    self.assertEqual(points_image.shape, (2,))
    self.assertAllClose(points_image, [cx, cy])

  def test_project_principal_point_batched(self):
    """3D points on +Z axis should project to principal point of the camera."""
    image_width, image_height = 640, 480
    fx, fy, cx, cy = 500.0, 500.0, image_width / 2, image_height / 2
    camera_model = cameras.PinholeCamera(
        K=tf.constant([[fx, 0., cx], [0., fy, cy], [0., 0., 1.]]),
        image_width=image_width,
        image_height=image_height)

    points_camera = np.tile(np.array([0., 0., 1.]), (10, 5, 1))
    points_image = camera_model.project(points_camera)
    self.assertEqual(points_image.shape, (10, 5, 2))
    self.assertAllClose(points_image, np.tile([cx, cy], (10, 5, 1)))

  def test_project_points(self):
    """Tests projection of random set of 3D points."""
    num_points = 10
    points_camera = np.random.rand(num_points, 3)
    intrinsics = np.array([[500, 0., 320.], [0., 500., 120.], [0., 0., 1.]])
    camera_model = cameras.PinholeCamera(
        K=intrinsics, image_width=640, image_height=240)
    points_image = camera_model.project(points_camera)
    self.assertEqual(points_image.shape, (num_points, 2))
    for point_camera, point_image in zip(points_camera, points_image):
      expected = intrinsics.dot(point_camera)
      expected = expected[:2] / expected[2]
      self.assertAllClose(point_image, expected)

  def test_unproject_principal_point_single(self):
    """Principal point should unproject to camera space points along +Z."""
    image_width, image_height = 640, 480
    fx, fy, cx, cy = 500.0, 500.0, image_width / 2, image_height / 2
    camera_model = cameras.PinholeCamera(
        K=np.array([[fx, 0., cx], [0., fy, cy], [0., 0., 1.]]),
        image_width=image_width,
        image_height=image_height)

    points_image = np.array([cx, cy])
    points_camera = camera_model.unproject(points_image)
    self.assertAllClose(points_camera, [0., 0., 1.])

  def test_unproject_principal_point_batched(self):
    """Principal point should unproject to camera space points along +Z."""
    image_width, image_height = 640, 480
    fx, fy, cx, cy = 500.0, 500.0, image_width / 2, image_height / 2
    camera_model = cameras.PinholeCamera(
        K=np.array([[fx, 0., cx], [0., fy, cy], [0., 0., 1.]]),
        image_width=image_width,
        image_height=image_height,
    )

    points_image = np.tile(np.array([cx, cy]), (6, 5, 4, 1))
    points_camera = camera_model.unproject(points_image)
    self.assertEqual(points_camera.shape, (6, 5, 4, 3))
    self.assertAllClose(points_camera,
                        np.tile(np.array([0., 0., 1.]), (6, 5, 4, 1)))

  def test_unproject_points_to_ray_directions(self):
    """Unproject a grid of 2D image points as rays and verify ray properties."""
    num_points = 10
    image_width, image_height = 640, 480
    points_image = np.column_stack(
        (np.random.uniform(0.0, image_width, num_points),
         np.random.uniform(0.0, image_height, num_points)))
    intrinsics = np.array([[500, 0., 320.], [0., 500., 120.], [0., 0., 1.]])
    camera_model = cameras.PinholeCamera(
        K=intrinsics, image_width=image_width, image_height=image_height)
    points_camera = camera_model.unproject(points_image, to_rays=True)
    self.assertEqual(points_camera.shape, (num_points, 3))

    # Assert thet each ray is a unit vector.
    self.assertAllClose(tf.norm(points_camera, axis=-1), tf.ones(num_points))

    # Assert thet each ray project backs to the same image coordinate.
    for point_camera, point_image in zip(points_camera, points_image):
      expected = intrinsics.dot(point_camera)
      expected = expected[:2] / expected[2]
      self.assertAllClose(point_image, expected)

  def test_unproject_points_to_unit_zplane(self):
    """Unproject a grid of 2D image points to z=1 plane and verify correctness."""
    num_points = 10
    image_width, image_height = 640, 480
    points_image = np.column_stack(
        (np.random.uniform(0.0, image_width, num_points),
         np.random.uniform(0.0, image_height, num_points)))
    intrinsics = np.array([[500, 0., 320.], [0., 500., 120.], [0., 0., 1.]])
    camera_model = cameras.PinholeCamera(
        K=intrinsics, image_width=image_width, image_height=image_height)
    points_camera = camera_model.unproject(points_image, to_rays=False)
    self.assertEqual(points_camera.shape, (num_points, 3))

    # Assert thet each ray end point is at z=1.
    self.assertAllClose(points_camera[:, 2], tf.ones(num_points))

    # Assert thet each ray project backs to the same image coordinate.
    for point_camera, point_image in zip(points_camera, points_image):
      expected = intrinsics.dot(point_camera)
      expected = expected[:2] / expected[2]
      self.assertAllClose(point_image, expected)

  def test_unproject_and_project(self):
    """Unproject image points to 3D and then project back to verify they are same."""
    # Create an pinhole camera.
    camera_model = cameras.PinholeCamera.from_intrinsics(
        image_size_in_pixels=(320, 240), focal_length_in_pixels=(400.0, 400.0))

    # Calling unproject without argument unprojects all pixel centers.
    ray_directions = camera_model.unproject()

    # Re-project the ray directions to verify they are same.
    projections = camera_model.project(ray_directions)
    self.assertAllClose(projections, camera_model.pixel_centers(), atol=1e-3)

  def test_pixel_centers(self):
    """Verifies the correctness of `pixel_centers` with an alternate implementation."""
    # Create an pinhole camera.
    image_width, image_height = 320, 240
    camera_model = cameras.PinholeCamera.from_intrinsics(
        image_size_in_pixels=(image_width, image_height),
        focal_length_in_pixels=(400.0, 400.0))

    # Create an array (H, W, 2) containing pixel center coordinates.
    pixel_centers = tf.meshgrid(
        tf.linspace(0.5, image_width - 0.5, image_width),
        tf.linspace(0.5, image_height - 0.5, image_height),
        indexing='xy')
    pixel_centers = tf.stack(pixel_centers, axis=-1)

    # Make sure the pixel centers are correct.
    self.assertAllClose(camera_model.pixel_centers(), pixel_centers)
