def ray_localization(
    heading: tuple,
    position: tuple,
    engine,
    use_heading_filter=True,
    return_on_lane=False,
) -> Union[List[Tuple], Tuple]:
    """
    Get the index of the lane closest to a physx_world position.
    Only used when smoething is on lane ! Otherwise fall back to use get_closest_lane()
    :param heading: heading to help filter lanes
    :param position: a physx_world position [m].
    :param engine: BaseEngine class
    :param return_all_result: return a list instead of the lane with min L1 distance
    :return: list(closest lane) or closest lane.
    """
    on_lane = False

    if len(position) == 3:
        position = position[:2]
    assert len(position) == 2
    assert len(heading) == 2

    results = engine.physics_world.static_world.rayTestAll(panda_vector(position, 1.0), panda_vector(position, -1.0))
    lane_index_dist = []
    if results.hasHits():
        for res in results.getHits():
            if MetaDriveType.is_lane(res.getNode().getName()):
                on_lane = True
                lane = get_object_from_node(res.getNode())
                long, _ = lane.local_coordinates(position)
                lane_heading = lane.heading_theta_at(long)

                # dir = np.array([math.cos(lane_heading), math.sin(lane_heading)])
                # dot_result = dir.dot(heading)

                dot_result = math.cos(lane_heading) * heading[0] + math.sin(lane_heading) * heading[1]
                cosangle = dot_result / (
                    norm(math.cos(lane_heading), math.sin(lane_heading)) * norm(heading[0], heading[1])
                )

                if use_heading_filter:
                    if cosangle > 0:
                        lane_index_dist.append((lane, lane.index, lane.distance(position)))
                else:
                    lane_index_dist.append((lane, lane.index, lane.distance(position)))
    # default return all result
    ret = []
    if len(lane_index_dist) > 0:
        ret = sorted(lane_index_dist, key=lambda k: k[2])

    # sorted(ret, key=lambda k: k[2]) what a stupid bug. sorted is not an inplace operation
    return (ret, on_lane) if return_on_lane else ret
    # else:
    #     if len(lane_index_dist) > 0:
    #         ret_index = np.argmin([d for _, _, d in lane_index_dist])
    #         lane, index, dist = lane_index_dist[ret_index]
    #     else:
    #         lane, index, dist = None, None, None
    #     return (lane, index) if not return_on_lane else (lane, index, on_lane)

    Get the index of the lane closest to a physx_world position.
    Only used when smoething is on lane ! Otherwise fall back to use get_closest_lane()
    :param heading: heading to help filter lanes
    :param position: a physx_world position [m].
    :param engine: BaseEngine class
    :param return_all_result: return a list instead of the lane with min L1 distance
    :return: list(closest lane) or closest lane.
    
metadrive/utils/pg/utils.py