'''
将点云投影到新的图像上

required:
    - numpy
'''

import numpy as np

def project_to_image(world_points: np.ndarray, w2c: np.ndarray, 
                      K: np.ndarray, width: int, height: int) -> np.ndarray:
    '''
    将点云投影到新的图像上，生成深度图

    Args:
        world_points: 世界坐标点云，形状为(N, 4)的齐次坐标
        w2c: 世界坐标到相机坐标变换矩阵 (4x4)
        K: 相机内参矩阵 (3x3)
        width: 图像宽度
        height: 图像高度

    Returns:
        深度图，值为无穷大表示没有投影到该像素的点
    '''
    # 检查输入数据的维度
    if world_points.shape[1] != 4:
        raise ValueError(f"世界点云应为齐次坐标 (N,4)，但得到 {world_points.shape}")
        
    # 初始化深度图
    depth_image = np.full((height, width), np.inf, dtype=np.float32)
    
    # 将世界坐标系点转换到相机坐标系
    points_camera = (w2c @ world_points.T).T
    
    # 筛选相机前方的点 (z > 0)
    valid_mask = points_camera[:, 2] > 0
    points_camera = points_camera[valid_mask]
    
    if len(points_camera) == 0:
        return depth_image
    
    # 投影到归一化图像平面 - 修复广播错误
    # 只使用深度值z进行归一化，而不是整个z之后的子数组
    z_values = points_camera[:, 2:3]  # 使用 2:3 保持列向量形式
    points_normalized = points_camera[:, :3] / z_values
    
    # 应用相机内参矩阵
    points_image = (K @ points_normalized[:, :3].T).T
    
    # 提取像素坐标和深度值
    x_pixels = np.round(points_image[:, 0]).astype(int)
    y_pixels = np.round(points_image[:, 1]).astype(int)
    depths = points_camera[:, 2]
    
    # 筛选出图像范围内的点
    inbounds = (x_pixels >= 0) & (x_pixels < width) & (y_pixels >= 0) & (y_pixels < height)
    x_pixels = x_pixels[inbounds]
    y_pixels = y_pixels[inbounds]
    depths = depths[inbounds]
    
    # 使用numpy的高级索引更新深度图 (更高效的实现)
    # 为每个像素找到最小深度值
    indices = y_pixels * width + x_pixels
    np.minimum.at(depth_image.reshape(-1), indices, depths)
    
    return depth_image