import open3d as o3d
import numpy as np
from scipy.spatial import KDTree
import copy


def compute_point_to_mesh_distance(point_cloud, mesh):
    """
    计算点云到 Mesh 的距离

    Args:
        point_cloud: Open3D 点云对象
        mesh: Open3D 网格对象

    Returns:
        distances: 每个点到最近三角形面的距离
    """
    # 方法1：使用 Open3D 内置函数（新版本）
    try:
        distances = mesh.compute_point_cloud_distance(point_cloud)
        return np.asarray(distances)
    except AttributeError:
        pass

    # 方法2：使用 mesh 顶点最近邻（近似）
    mesh_points = np.asarray(mesh.vertices)
    points = np.asarray(point_cloud.points)

    tree = KDTree(mesh_points)
    distances, indices = tree.query(points)

    return distances


def compute_statistics(distances):
    """计算距离统计信息"""
    stats = {
        'mean': np.mean(distances),
        'rmse': np.sqrt(np.mean(distances ** 2)),
        'max': np.max(distances),
        'min': np.min(distances),
        'std': np.std(distances),
        'median': np.median(distances),
    }

    # 分位数
    percentiles = [50, 75, 90, 95, 99]
    for p in percentiles:
        stats[f'percentile_{p}'] = np.percentile(distances, p)

    return stats


def visualize_distance(point_cloud, mesh, distances, max_dist=None):
    """可视化距离（按距离着色）"""
    pcd_colored = copy.deepcopy(point_cloud)

    if max_dist is None:
        max_dist = np.percentile(distances, 95)

    # 生成颜色（红：远，蓝：近）
    colors = np.zeros((len(distances), 3))
    norm_dist = np.clip(distances / (max_dist + 1e-8), 0, 1)
    colors[:, 0] = norm_dist  # 红色通道（远）
    colors[:, 2] = 1 - norm_dist  # 蓝色通道（近）

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

    # Mesh 着色为灰色
    mesh_colored = copy.deepcopy(mesh)
    mesh_colored.paint_uniform_color([0.5, 0.5, 0.5])

    # 可视化
    o3d.visualization.draw_geometries(
        [mesh_colored, pcd_colored],
        window_name="Distance to Mesh (Blue=Near, Red=Far)",
        width=1024,
        height=768
    )


def print_statistics(stats):
    """打印统计信息"""
    print("=" * 60)
    print("点云到 Mesh 距离统计")
    print("=" * 60)
    print(f"平均距离: {stats['mean']:.6f} m ({stats['mean'] * 1000:.3f} mm)")
    print(f"RMSE:     {stats['rmse']:.6f} m ({stats['rmse'] * 1000:.3f} mm)")
    print(f"中位数:   {stats['median']:.6f} m ({stats['median'] * 1000:.3f} mm)")
    print(f"标准差:   {stats['std']:.6f} m ({stats['std'] * 1000:.3f} mm)")
    print(f"最大距离: {stats['max']:.6f} m ({stats['max'] * 1000:.3f} mm)")
    print(f"最小距离: {stats['min']:.6f} m ({stats['min'] * 1000:.3f} mm)")
    print("\n距离分位数:")
    print(f"  50%:  {stats['percentile_50']:.6f} m ({stats['percentile_50'] * 1000:.3f} mm)")
    print(f"  75%:  {stats['percentile_75']:.6f} m ({stats['percentile_75'] * 1000:.3f} mm)")
    print(f"  90%:  {stats['percentile_90']:.6f} m ({stats['percentile_90'] * 1000:.3f} mm)")
    print(f"  95%:  {stats['percentile_95']:.6f} m ({stats['percentile_95'] * 1000:.3f} mm)")
    print(f"  99%:  {stats['percentile_99']:.6f} m ({stats['percentile_99'] * 1000:.3f} mm)")


# ========== 主程序 ==========
if __name__ == "__main__":
    # 读取文件（使用原始字符串避免转义问题）
    file_path = r"E:\point_distance\data\40_noisy.ply"
    mesh_path = r"E:\point_distance\data\40.stl"

    print("读取点云...")
    source_pcd = o3d.io.read_point_cloud(file_path)

    print("读取 Mesh...")
    mesh = o3d.io.read_triangle_mesh(mesh_path)

    # 检查读取是否成功
    if len(source_pcd.points) == 0:
        print("错误：点云读取失败，请检查文件路径")
        exit()

    if len(mesh.vertices) == 0:
        print("错误：Mesh 读取失败，请检查文件路径")
        exit()

    print(f"\n点云点数: {len(source_pcd.points)}")
    print(f"Mesh 顶点数: {len(mesh.vertices)}")
    print(f"Mesh 面数: {len(mesh.triangles)}")

    # 可选：Mesh 法线计算（用于更好的可视化）
    mesh.compute_vertex_normals()

    # 计算距离
    print("\n计算点云到 Mesh 的距离...")
    distances = compute_point_to_mesh_distance(source_pcd, mesh)

    # 计算统计信息
    stats = compute_statistics(distances)

    # 打印统计信息
    print_statistics(stats)

    # 可视化
    print("\n可视化...")
    visualize_distance(source_pcd, mesh, distances)