import torch


def compute_bounds_from_intersect_points(rays_o, intersect_indices,
                                         intersect_points):
    """Computes bounds from intersection points.

    Note: Make sure that inputs are in the same coordiante frame.

    Args:
        rays_o: [R, 3] float tensor
        intersect_indices: [R', 1] float tensor
        intersect_points: [R', 2, 3] float tensor

    Returns:
        intersect_bounds: [R', 2] float tensor

    where R is the number of rays and R' is the number of intersecting rays.
    """
    intersect_rays_o = rays_o[intersect_indices]  # [R', 1, 3]
    intersect_diff = intersect_points - intersect_rays_o  # [R', 2, 3]
    intersect_bounds = torch.norm(intersect_diff, dim=2)  # [R', 2]

    # Sort the bounds so that near comes before far for all rays.
    intersect_bounds, _ = torch.sort(intersect_bounds, dim=1)  # [R', 2]

    # For some reason the sort function returns [R', ?] instead of [R', 2], so we
    # will explicitly reshape it.
    intersect_bounds = intersect_bounds.view(-1, 2)  # [R', 2]
    return intersect_bounds