import open3d as o3d
import numpy as np
import argparse
import os


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


def visualize_curvature(pcd, curvatures, color_scale=3.0):
    """
    根据曲率可视化点云

    Args:
        pcd: Open3D 点云对象
        curvatures: 曲率数组
        color_scale: 颜色缩放因子
    """
    colors = np.zeros((len(pcd.points), 3))
    # 红色通道表示曲率（高曲率区域为红色）
    colors[:, 0] = np.clip(curvatures * color_scale, 0, 1)
    # 绿色通道为互补色（低曲率区域为绿色）
    colors[:, 1] = np.clip(1 - curvatures * color_scale, 0, 1)
    # 蓝色通道设为0
    colors[:, 2] = 0

    pcd.colors = o3d.utility.Vector3dVector(colors)


def main():
    parser = argparse.ArgumentParser(description='计算点云曲率并可视化')
    parser.add_argument('--ply_path', type=str,default="./High_Speed_Data/PointCloud/11.ply", help='输入点云文件路径 (.ply)')
    parser.add_argument('--radius', type=float, default=2, help='搜索半径 (默认: 50)')
    parser.add_argument('--max_nn', type=int, default=30, help='最大邻居数 (默认: 30)')
    parser.add_argument('--color_scale', type=float, default=3.0, help='颜色缩放因子 (默认: 3.0)')
    parser.add_argument('--save', type=str, default=None, help='保存可视化图像路径 (可选)')

    args = parser.parse_args()

    # 检查文件是否存在
    if not os.path.exists(args.ply_path):
        print(f"错误: 文件不存在 - {args.ply_path}")
        return

    print(f"读取点云: {args.ply_path}")
    pcd = o3d.io.read_point_cloud(args.ply_path)
    # points=np.array(pcd.points)
    if len(pcd.points) == 0:
        print("错误: 读取点云失败，文件可能为空或格式不正确")
        return

    print(f"点云点数: {len(pcd.points)}")

    # 计算曲率
    print("计算曲率中...")
    curvatures = compute_curvature_open3d(pcd, radius=args.radius, max_nn=args.max_nn)

    # 统计曲率信息
    print(f"曲率统计:")
    print(f"  最小值: {curvatures.min():.6f}")
    print(f"  最大值: {curvatures.max():.6f}")
    print(f"  平均值: {curvatures.mean():.6f}")
    print(f"  标准差: {curvatures.std():.6f}")

    # 可视化
    print("可视化曲率...")
    visualize_curvature(pcd, curvatures, color_scale=args.color_scale)

    # 保存截图（如果指定）
    if args.save:
        o3d.io.write_point_cloud(args.save, pcd)
        print(f"点云已保存: {args.save}")

    # 显示窗口
    o3d.visualization.draw_geometries([pcd], window_name=f"Curvature Visualization - {os.path.basename(args.ply_path)}")


if __name__ == "__main__":
    main()