# Path: sunds/core/tf_geometry/cameras.py
# class AbstractCamera(abc.ABC):
# class PinholeCamera(AbstractCamera):
#   def project(self, points: TensorLike) -> tf.Tensor:
#   def unproject(self,
#                 points: Optional[TensorLike] = None,
#                 to_rays: bool = True) -> tf.Tensor:
#   def pixel_centers(self) -> tf.Tensor:
#   def __post_init__(self):
#   def from_fov(
#       cls,
#       *,
#       image_width: int,
#       image_height: int,
#       horizontal_fov_in_degrees: float,
#       vertical_fov_in_degrees: Optional[float] = None) -> 'PinholeCamera':
#     def focal_length_from_fov(image_size, fov_in_degrees):
#   def from_intrinsics(
#       cls,
#       *,
#       image_size_in_pixels: Tuple[int, int],
#       focal_length_in_pixels: Tuple[float, float],
#       principal_point_in_pixels: Optional[Tuple[float, float]] = None,
#       skew: float = 0.0,
#   ) -> 'PinholeCamera':
#   def project(self, points: TensorLike) -> tf.Tensor:
#   def unproject(self,
#                 points: Optional[TensorLike] = None,
#                 to_rays: bool = True) -> tf.Tensor:
#   def pixel_centers(self) -> tf.Tensor:
#   K: TensorLike  # pylint: disable=invalid-name
# 
# Path: sunds/core/tf_geometry/isometry.py
# class Isometry:
#   R: TensorLike  # pylint: disable=invalid-name
#   def __post_init__(self):
#   def from_matrix(cls, matrix: TensorLike) -> 'Isometry':
#   def matrix3x4(self) -> tf.Tensor:
#   def matrix4x4(self) -> tf.Tensor:
#   def inverse(self) -> 'Isometry':
#   def compose(self, other: 'Isometry') -> 'Isometry':
#   def transform_points(self, points: TensorLike) -> tf.Tensor:
#   def __mul__(
#       self, other: Union['Isometry',
#                          TensorLike]) -> Union['Isometry', tf.Tensor]:
# 
# Path: sunds/core/tf_geometry/rays.py
# def rays_from_image_points(camera: cameras.CameraType,
#                            world_from_camera: isometry.Isometry,
#                            points_image: TensorLike):
# def rays_from_image_grid(camera: cameras.CameraType,
#                          world_from_camera: isometry.Isometry):
# def _rays_from_directions(ray_directions: TensorLike,
#                           world_from_camera: isometry.Isometry):
# def depth_samples_along_rays(near_depth: TensorLike,
#                              far_depth: TensorLike,
#                              num_samples_per_ray: int,
#                              method: str = 'linspace_depth') -> tf.Tensor:
# def jitter_depth_samples_along_rays(point_depths: TensorLike) -> tf.Tensor:
# def point_samples_along_rays(ray_origins: TensorLike,
#                              ray_directions: TensorLike,
#                              point_depths: TensorLike) -> tf.Tensor:


from absl.testing import parameterized
from sunds.core.tf_geometry import cameras
from sunds.core.tf_geometry import isometry
from sunds.core.tf_geometry import rays
import numpy as np
import tensorflow as tf
    # Generate rays.
    ray_origins, ray_directions = rays.rays_from_image_grid(
        camera=camera, world_from_camera=world_from_camera)

    # Generate ray end points in world frame and project back to camera.
    points_world = ray_origins + 5.0 * ray_directions
    points_camera = world_from_camera.inverse() * points_world
    points_image = camera.project(points_camera)

    # Create an array (H, W, 2) containing pixel center coordinates.
    pixel_centers = np.stack(
        np.meshgrid(
            range(camera.image_width),
            range(camera.image_height),
            indexing='xy'),
        axis=-1)
    pixel_centers = pixel_centers.astype(np.float32) + 0.5

    self.assertAllClose(points_image, pixel_centers, atol=1e-03)


class SamplesAlongRaysTest(parameterized.TestCase, tf.test.TestCase):

  @parameterized.named_parameters(
      ('method=linspace_depth', 'linspace_depth'),
      ('method=linspace_disparity', 'linspace_disparity'),
      ('method=geomspace_depth', 'geomspace_depth'),
  )
  def test_depth_samples_scalar_range(self, method):
    """Checks depth samples with a single scalar input depth range."""

# Path: sunds/core/tf_geometry/cameras.py
# class AbstractCamera(abc.ABC):
# class PinholeCamera(AbstractCamera):
#   def project(self, points: TensorLike) -> tf.Tensor:
#   def unproject(self,
#                 points: Optional[TensorLike] = None,
#                 to_rays: bool = True) -> tf.Tensor:
#   def pixel_centers(self) -> tf.Tensor:
#   def __post_init__(self):
#   def from_fov(
#       cls,
#       *,
#       image_width: int,
#       image_height: int,
#       horizontal_fov_in_degrees: float,
#       vertical_fov_in_degrees: Optional[float] = None) -> 'PinholeCamera':
#     def focal_length_from_fov(image_size, fov_in_degrees):
#   def from_intrinsics(
#       cls,
#       *,
#       image_size_in_pixels: Tuple[int, int],
#       focal_length_in_pixels: Tuple[float, float],
#       principal_point_in_pixels: Optional[Tuple[float, float]] = None,
#       skew: float = 0.0,
#   ) -> 'PinholeCamera':
#   def project(self, points: TensorLike) -> tf.Tensor:
#   def unproject(self,
#                 points: Optional[TensorLike] = None,
#                 to_rays: bool = True) -> tf.Tensor:
#   def pixel_centers(self) -> tf.Tensor:
#   K: TensorLike  # pylint: disable=invalid-name
# 
# Path: sunds/core/tf_geometry/isometry.py
# class Isometry:
#   R: TensorLike  # pylint: disable=invalid-name
#   def __post_init__(self):
#   def from_matrix(cls, matrix: TensorLike) -> 'Isometry':
#   def matrix3x4(self) -> tf.Tensor:
#   def matrix4x4(self) -> tf.Tensor:
#   def inverse(self) -> 'Isometry':
#   def compose(self, other: 'Isometry') -> 'Isometry':
#   def transform_points(self, points: TensorLike) -> tf.Tensor:
#   def __mul__(
#       self, other: Union['Isometry',
#                          TensorLike]) -> Union['Isometry', tf.Tensor]:
# 
# Path: sunds/core/tf_geometry/rays.py
# def rays_from_image_points(camera: cameras.CameraType,
#                            world_from_camera: isometry.Isometry,
#                            points_image: TensorLike):
# def rays_from_image_grid(camera: cameras.CameraType,
#                          world_from_camera: isometry.Isometry):
# def _rays_from_directions(ray_directions: TensorLike,
#                           world_from_camera: isometry.Isometry):
# def depth_samples_along_rays(near_depth: TensorLike,
#                              far_depth: TensorLike,
#                              num_samples_per_ray: int,
#                              method: str = 'linspace_depth') -> tf.Tensor:
# def jitter_depth_samples_along_rays(point_depths: TensorLike) -> tf.Tensor:
# def point_samples_along_rays(ray_origins: TensorLike,
#                              ray_directions: TensorLike,
#                              point_depths: TensorLike) -> tf.Tensor:


# Path: sunds/core/tf_geometry/rays_test.py
from absl.testing import parameterized
from sunds.core.tf_geometry import cameras
from sunds.core.tf_geometry import isometry
from sunds.core.tf_geometry import rays
import numpy as np
import tensorflow as tf

    # Generate rays.
    ray_origins, ray_directions = rays.rays_from_image_grid(
        camera=camera, world_from_camera=world_from_camera)

    # Generate ray end points in world frame and project back to camera.
    points_world = ray_origins + 5.0 * ray_directions
    points_camera = world_from_camera.inverse() * points_world
    points_image = camera.project(points_camera)

    # Create an array (H, W, 2) containing pixel center coordinates.
    pixel_centers = np.stack(
        np.meshgrid(
            range(camera.image_width),
            range(camera.image_height),
            indexing='xy'),
        axis=-1)
    pixel_centers = pixel_centers.astype(np.float32) + 0.5

    self.assertAllClose(points_image, pixel_centers, atol=1e-03)


class SamplesAlongRaysTest(parameterized.TestCase, tf.test.TestCase):

  @parameterized.named_parameters(
      ('method=linspace_depth', 'linspace_depth'),
      ('method=linspace_disparity', 'linspace_disparity'),
      ('method=geomspace_depth', 'geomspace_depth'),
  )
  def test_depth_samples_scalar_range(self, method):
    """Checks depth samples with a single scalar input depth range."""

    point_depths = rays.depth_samples_along_rays(