# Path: ldif/util/geom_util.py
# def ray_sphere_intersect(ray_start, ray_direction, sphere_center, sphere_radius,
#                          max_t):
# def to_homogeneous(t, is_point):
# def transform_points_with_normals(points, tx, normals=None):
# def transform_featured_points(points, tx):
# def rotation_to_tx(rot):
# def extract_points_near_origin(points, count, features=None):
# def local_views_of_shape(global_points,
#                          world2local,
#                          local_point_count,
#                          global_normals=None,
#                          global_features=None,
#                          zeros_invalid=False,
#                          zero_threshold=1e-6,
#                          expand_region=True,
#                          threshold=4.0):
# def chamfer_distance(pred, target):
# def dodeca_parameters(dodeca_idx):
# def get_camera_to_world(viewpoint, center, world_up):
# def get_dodeca_camera_to_worlds():
# def gaps_depth_render_to_xyz(model_config, depth_image, camera_parameters):
# def angle_of_rotation_to_2d_rotation_matrix(angle_of_rotation):
# def fractional_vector_projection(e0, e1, p, falloff=2.0):
# def rotate_about_point(angle_of_rotation, point, to_rotate):
# def interpolate_from_grid(samples, grid):
# def reflect(samples, reflect_x=False, reflect_y=False, reflect_z=False):
# def z_reflect(samples):
# def get_world_to_camera(idx):
# def transform_depth_dodeca_to_xyz_dodeca(depth_dodeca):
# def transform_depth_dodeca_to_xyz_dodeca_np(depth_dodeca):
# def _unbatch(arr):
# def to_homogenous_np(arr, is_point=True):
# def depth_to_cam_np(im, xfov=0.5):
# def apply_tx_np(samples, tx, is_point=True):
# def depth_image_to_sdf_constraints(im, cam2world, xfov=0.5):
# def depth_dodeca_to_sdf_constraints(depth_ims):
# def depth_dodeca_to_samples(dodeca):
# def depth_image_to_class_constraints(im, cam2world, xfov=0.5):
# def depth_image_to_samples(im, cam2world, xfov=0.5):  # pylint:disable=unused-argument
# def apply_4x4(tensor, tx, are_points=True, batch_rank=None, sample_rank=None):
# def depth_image_to_xyz_image(depth_images, world_to_camera, xfov=0.5):
# def interpolate_from_grid_coordinates(samples, grid):
# 
# Path: ldif/util/np_util.py
# def batch_np(arr, batch_size):
# def make_coordinate_grid(height, width, is_screen_space, is_homogeneous):
# def make_coordinate_grid_3d(length, height, width, is_screen_space,
#                             is_homogeneous):
# def filter_valid(mask, vals):
# def zero_by_mask(mask, vals, replace_with=0.0):
# def make_mask(im, thresh=0.0):
# def make_pixel_mask(im):
# def thresh_and_radius_to_distance(radius, thresh):
# def plot_rbfs_at_thresh(centers, radii, thresh=0.5):
# def plot_rbfs(centers, radii, scale=10.0):
# def cube_and_render(volume, thresh):
# def sample_surface(quadrics, centers, radii, length, height, width,
#                    renormalize):
# def visualize_prediction(quadrics,
#                          centers,
#                          radii,
#                          renormalize,
#                          thresh=0.0,
#                          input_volumes=None):
# 
# Path: ldif/util/tf_util.py
# def assert_shape(tensor, shape, name):
# def log(msg, t):
# def tile_new_axis(t, axis, length):
# def zero_by_mask(mask, vals, replace_with=0.0):
# def remove_element(t, elt, axis):


import importlib
import tensorflow as tf
from ldif.util import geom_util
from ldif.util import np_util
from ldif.util import tf_util
# pylint: enable=g-bad-import-order

importlib.reload(tf_util)


def gaps_depth_image_to_cam_image(depth_image, xfov):
  """Converts a GAPS depth image tensor to a camera-space image tensor.

  Args:
    depth_image: Tensor with shape [batch_size, height, width, 1].
    xfov: Scalar or tensor with shape [1] or [batch_size].

  Returns:
    cam_image: Tensor with shape [batch_size, height, width, 3].
  """
  batch_size, height, width = depth_image.get_shape().as_list()[:3]
  depth_image = tf.ensure_shape(depth_image, [batch_size, height, width, 1])
  if isinstance(xfov, float):
    xfov = tf.constant([xfov], dtype=tf.float32)
    xfov = tf.tile(xfov, [batch_size])
  else:
    xfov = tf.reshape(xfov, [batch_size])
#  if xfov.get_shape().as_list()[0] == 1:
#    xfov = tf.tile(xfov, [batch_size])
#  else:
#    assert xfov.get_shape().as_list()[0] == batch_size

  pixel_coords = np_util.make_coordinate_grid(
      height, width, is_screen_space=False, is_homogeneous=False)
  # Values should go from -1 -> 1, not from 0 -> 1:

# Path: ldif/util/geom_util.py
# def ray_sphere_intersect(ray_start, ray_direction, sphere_center, sphere_radius,
#                          max_t):
# def to_homogeneous(t, is_point):
# def transform_points_with_normals(points, tx, normals=None):
# def transform_featured_points(points, tx):
# def rotation_to_tx(rot):
# def extract_points_near_origin(points, count, features=None):
# def local_views_of_shape(global_points,
#                          world2local,
#                          local_point_count,
#                          global_normals=None,
#                          global_features=None,
#                          zeros_invalid=False,
#                          zero_threshold=1e-6,
#                          expand_region=True,
#                          threshold=4.0):
# def chamfer_distance(pred, target):
# def dodeca_parameters(dodeca_idx):
# def get_camera_to_world(viewpoint, center, world_up):
# def get_dodeca_camera_to_worlds():
# def gaps_depth_render_to_xyz(model_config, depth_image, camera_parameters):
# def angle_of_rotation_to_2d_rotation_matrix(angle_of_rotation):
# def fractional_vector_projection(e0, e1, p, falloff=2.0):
# def rotate_about_point(angle_of_rotation, point, to_rotate):
# def interpolate_from_grid(samples, grid):
# def reflect(samples, reflect_x=False, reflect_y=False, reflect_z=False):
# def z_reflect(samples):
# def get_world_to_camera(idx):
# def transform_depth_dodeca_to_xyz_dodeca(depth_dodeca):
# def transform_depth_dodeca_to_xyz_dodeca_np(depth_dodeca):
# def _unbatch(arr):
# def to_homogenous_np(arr, is_point=True):
# def depth_to_cam_np(im, xfov=0.5):
# def apply_tx_np(samples, tx, is_point=True):
# def depth_image_to_sdf_constraints(im, cam2world, xfov=0.5):
# def depth_dodeca_to_sdf_constraints(depth_ims):
# def depth_dodeca_to_samples(dodeca):
# def depth_image_to_class_constraints(im, cam2world, xfov=0.5):
# def depth_image_to_samples(im, cam2world, xfov=0.5):  # pylint:disable=unused-argument
# def apply_4x4(tensor, tx, are_points=True, batch_rank=None, sample_rank=None):
# def depth_image_to_xyz_image(depth_images, world_to_camera, xfov=0.5):
# def interpolate_from_grid_coordinates(samples, grid):
# 
# Path: ldif/util/np_util.py
# def batch_np(arr, batch_size):
# def make_coordinate_grid(height, width, is_screen_space, is_homogeneous):
# def make_coordinate_grid_3d(length, height, width, is_screen_space,
#                             is_homogeneous):
# def filter_valid(mask, vals):
# def zero_by_mask(mask, vals, replace_with=0.0):
# def make_mask(im, thresh=0.0):
# def make_pixel_mask(im):
# def thresh_and_radius_to_distance(radius, thresh):
# def plot_rbfs_at_thresh(centers, radii, thresh=0.5):
# def plot_rbfs(centers, radii, scale=10.0):
# def cube_and_render(volume, thresh):
# def sample_surface(quadrics, centers, radii, length, height, width,
#                    renormalize):
# def visualize_prediction(quadrics,
#                          centers,
#                          radii,
#                          renormalize,
#                          thresh=0.0,
#                          input_volumes=None):
# 
# Path: ldif/util/tf_util.py
# def assert_shape(tensor, shape, name):
# def log(msg, t):
# def tile_new_axis(t, axis, length):
# def zero_by_mask(mask, vals, replace_with=0.0):
# def remove_element(t, elt, axis):


# Path: ldif/util/gaps_util_tf.py
import importlib
import tensorflow as tf
from ldif.util import geom_util
from ldif.util import np_util
from ldif.util import tf_util

# pylint: enable=g-bad-import-order

importlib.reload(tf_util)


def gaps_depth_image_to_cam_image(depth_image, xfov):
  """Converts a GAPS depth image tensor to a camera-space image tensor.

  Args:
    depth_image: Tensor with shape [batch_size, height, width, 1].
    xfov: Scalar or tensor with shape [1] or [batch_size].

  Returns:
    cam_image: Tensor with shape [batch_size, height, width, 3].
  """
  batch_size, height, width = depth_image.get_shape().as_list()[:3]
  depth_image = tf.ensure_shape(depth_image, [batch_size, height, width, 1])
  if isinstance(xfov, float):
    xfov = tf.constant([xfov], dtype=tf.float32)
    xfov = tf.tile(xfov, [batch_size])
  else:
    xfov = tf.reshape(xfov, [batch_size])
#  if xfov.get_shape().as_list()[0] == 1:
#    xfov = tf.tile(xfov, [batch_size])
#  else:
#    assert xfov.get_shape().as_list()[0] == batch_size

  pixel_coords = np_util.make_coordinate_grid(
      height, width, is_screen_space=False, is_homogeneous=False)
  # Values should go from -1 -> 1, not from 0 -> 1:

  nic_x = np_util.batch_np(2 * pixel_coords[:, :, 0:1] - 1.0, batch_size)