import os
import numpy as np
import open3d as o3d
from glob import glob
from tifffile import imwrite as tif_write

# 相机内参矩阵
K = np.array([
    [1.30449e3, 0,         5.2602e2],
    [0,         1.30449e3, 1.07432e3],
    [0,         0,         1]
])
fx, fy = K[0, 0], K[1, 1]
cx, cy = K[0, 2], K[1, 2]

# 图像尺寸
height, width = 2000, 2000

def compute_uv_mask(points):
    """
    计算 u, v 像素坐标和有效 mask
    """
    X, Y, Z = points[:, 0], points[:, 1], points[:, 2]
    valid = Z > 0
    X, Y, Z = X[valid], Y[valid], Z[valid]

    u = ((X * fx) / Z + cx).astype(np.int32)
    v = ((Y * fy) / Z + cy).astype(np.int32)

    mask = (u >= 0) & (u < width) & (v >= 0) & (v < height)

    return u[mask], v[mask], Z[mask], valid, mask

def build_rgb_image(points, u, v, valid_mask):
    r, g, b = points[:, 3], points[:, 4], points[:, 5]
    r, g, b = r[valid_mask], g[valid_mask], b[valid_mask]

    color_image = np.zeros((height, width, 3), dtype=np.uint8)
    color_image[v, u, 0] = (r * 255).astype(np.uint8)
    color_image[v, u, 1] = (g * 255).astype(np.uint8)
    color_image[v, u, 2] = (b * 255).astype(np.uint8)
    return color_image

def build_depth_map(u, v, Z):
    depth_map = np.zeros((height, width), dtype=np.float32)
    depth_map[v, u] = Z
    return depth_map

def process_pcd_file(pcd_path, rgbd_dir):
    filename = os.path.splitext(os.path.basename(pcd_path))[0]
    print(f"Processing: {filename}")

    pcd = o3d.io.read_point_cloud(pcd_path)
    points = np.asarray(pcd.points)
    colors = np.asarray(pcd.colors)

    if points.shape[0] == 0 or colors.shape[0] == 0:
        print(f"Skipping empty or invalid PCD: {filename}")
        return

    points_colored = np.hstack((points, colors))  # Nx6

    # 一次性计算 u, v 和 mask
    u, v, Z, valid_z_mask, final_mask = compute_uv_mask(points_colored)

    # 生成 RGB 和 depth 图像
    color_img = build_rgb_image(points_colored[valid_z_mask], u, v, final_mask)
    depth_map = build_depth_map(u, v, Z)

    # 合并为 RGBD
    rgbd = np.zeros((height, width, 4), dtype=np.float32)
    rgbd[:, :, 0:3] = color_img.astype(np.float32)
    rgbd[:, :, 3] = depth_map

    # 保存 tiff
    tif_write(os.path.join(rgbd_dir, f"{filename}_rgbd.tiff"), rgbd, photometric='rgb')
    print(f"Saved RGBD TIFF for {filename}")

if __name__ == "__main__":
    input_dir = r"\\192.168.50.222\share2\zyh\data\rgbd_line\dianyun0731"
    output_base = os.path.dirname(input_dir)
    rgbd_dir = os.path.join(output_base, "rgbd_tiff")
    os.makedirs(rgbd_dir, exist_ok=True)

    pcd_files = glob(os.path.join(input_dir, "*.pcd"))
    print(f"Found {len(pcd_files)} PCD files.")

    for pcd_path in pcd_files:
        process_pcd_file(pcd_path, rgbd_dir)