import open3d as o3d
import numpy as np


def compute_curvature_open3d(pcd, radius=0.05, max_nn=30):
    """
    使用 Open3D 计算点云曲率

    Args:
        pcd: Open3D 点云对象
        radius: 搜索半径
        max_nn: 最大邻居数

    Returns:
        curvatures: (N,) 曲率数组
    """
    # 估计法线（同时获取协方差矩阵信息）
    pcd.estimate_normals(
        search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=max_nn)
    )

    # 获取点云数据
    points = np.asarray(pcd.points)
    n_points = len(points)
    curvatures = np.zeros(n_points)

    # 构建 KD-Tree
    tree = o3d.geometry.KDTreeFlann(pcd)

    for i in range(n_points):
        # 搜索邻域点
        [k, idx, _] = tree.search_radius_vector_3d(pcd.points[i], radius)

        if k < 3:
            curvatures[i] = 1.0
            continue

        # 获取邻域点坐标
        neighbors = points[idx]

        # 计算协方差矩阵
        centroid = np.mean(neighbors, axis=0)
        centered = neighbors - centroid
        cov = np.cov(centered.T)

        # 特征值分解
        eigenvalues = np.linalg.eigvalsh(cov)
        eigenvalues = np.sort(eigenvalues)

        # 计算曲率
        curvatures[i] = eigenvalues[0] / (eigenvalues.sum() + 1e-8)

    return curvatures


# 使用示例
pcd = o3d.io.read_point_cloud("E:\point_distance\data\PointCloud\\38.ply")
curvatures = compute_curvature_open3d(pcd, radius=50, max_nn=30)

# 将曲率作为颜色可视化
colors = np.zeros((len(pcd.points), 3))
colors[:, 0] = curvatures*3  # 红色通道表示曲率
colors[:, 1] = 1 - curvatures*3  # 绿色通道为互补色
pcd.colors = o3d.utility.Vector3dVector(colors)
o3d.visualization.draw_geometries([pcd])