import open3d as o3d
import numpy as np
from scipy.spatial import KDTree
import numpy as np
from scipy.spatial import KDTree
def chamfer_distance(pcd1, pcd2):
    """
    计算两个点云之间的 CD 距离（Chamfer Distance）

    Args:
        pcd1: Open3D 点云对象1
        pcd2: Open3D 点云对象2

    Returns:
        cd: Chamfer Distance 值
        cd1: 点云1 到 点云2 的平均距离
        cd2: 点云2 到 点云1 的平均距离
    """
    # 获取点坐标
    points1 = np.asarray(pcd1.points)
    points2 = np.asarray(pcd2.points)

    # 构建 KD-Tree 加速最近邻搜索
    tree1 = KDTree(points1)
    tree2 = KDTree(points2)

    # 计算 points1 中每个点到 points2 的最近距离
    dist1, _ = tree2.query(points1, k=1)
    cd1 = np.mean(dist1 ** 2)  # 平方距离的平均值

    # 计算 points2 中每个点到 points1 的最近距离
    dist2, _ = tree1.query(points2, k=1)
    cd2 = np.mean(dist2 ** 2)  # 平方距离的平均值

    # 总 CD 距离（双向平均）
    cd = cd1 + cd2

    return np.sqrt(cd), cd1, cd2

def bilateral_filter(pcd, sigma_s=0.01, sigma_r=0.1):
    """
    对点云进行双边滤波
    Args:
        pcd: Open3D点云对象
        sigma_s: 空间标准差
        sigma_r: 特征（法线）标准差
    """
    points = np.asarray(pcd.points)
    normals = np.asarray(pcd.normals)

    # 构建KD树
    tree = KDTree(points)

    filtered_points = np.zeros_like(points)

    for i, p in enumerate(points):
        # 搜索邻域点
        indices = tree.query_ball_point(p, r=2 * sigma_s)
        if len(indices) <= 1:
            filtered_points[i] = p
            continue

        neigh_points = points[indices]
        neigh_normals = normals[indices]

        # 计算空间权重
        spatial_dist = np.linalg.norm(neigh_points - p, axis=1)
        w_s = np.exp(-spatial_dist ** 2 / (2 * sigma_s ** 2))

        # 计算特征（法线）权重
        normal_diff = np.linalg.norm(neigh_normals - normals[i], axis=1)
        w_r = np.exp(-normal_diff ** 2 / (2 * sigma_r ** 2))

        # 双边权重
        w = w_s * w_r
        w = w / w.sum()  # 归一化

        # 加权平均
        filtered_points[i] = np.sum(neigh_points.T * w, axis=1)

    # 更新点云
    filtered_pcd = o3d.geometry.PointCloud()
    filtered_pcd.points = o3d.utility.Vector3dVector(filtered_points)
    filtered_pcd.normals = o3d.utility.Vector3dVector(normals)

    return filtered_pcd


# 使用示例
pcd = o3d.io.read_point_cloud("E:\point_distance\data\part_car_head.ply")
pcd.estimate_normals()
pcd_filtered = bilateral_filter(pcd, sigma_s=10, sigma_r=10)

cd=chamfer_distance(pcd,pcd_filtered)
print(cd)