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


def compute_keep_ratio_and_density(pcd, filtered_pcd, sample_size=1000):
    """
    计算直通滤波的保留率和密度变化

    Args:
        pcd: 原始点云
        filtered_pcd: 滤波后点云
        sample_size: 密度计算采样点数（加速计算）

    Returns:
        keep_ratio: 保留率 (%)
        density_change: 密度变化 (%)
    """
    n_orig = len(pcd.points)
    n_filt = len(filtered_pcd.points)

    # 1. 保留率
    keep_ratio = n_filt / n_orig * 100

    # 2. 密度变化（基于平均最近邻距离）
    points_orig = np.asarray(pcd.points)
    points_filt = np.asarray(filtered_pcd.points)

    # 采样（点数太多时加速）
    if len(points_orig) > sample_size:
        idx_orig = np.random.choice(len(points_orig), sample_size, replace=False)
        points_orig_sample = points_orig[idx_orig]
    else:
        points_orig_sample = points_orig

    if len(points_filt) > sample_size:
        idx_filt = np.random.choice(len(points_filt), sample_size, replace=False)
        points_filt_sample = points_filt[idx_filt]
    else:
        points_filt_sample = points_filt

    # 计算平均最近邻距离（密度指标）
    if len(points_orig_sample) > 5:
        tree_orig = KDTree(points_orig_sample)
        dist_orig, _ = tree_orig.query(points_orig_sample, k=2)  # k=2 排除自身
        avg_dist_orig = np.mean(dist_orig[:, 1])
    else:
        avg_dist_orig = 0

    if len(points_filt_sample) > 5:
        tree_filt = KDTree(points_filt_sample)
        dist_filt, _ = tree_filt.query(points_filt_sample, k=2)
        avg_dist_filt = np.mean(dist_filt[:, 1])
    else:
        avg_dist_filt = 0

    # 密度变化（距离增大 = 密度减小）
    if avg_dist_orig > 0:
        density_change = (avg_dist_filt - avg_dist_orig) / avg_dist_orig * 100
    else:
        density_change = 0

    return keep_ratio, density_change


# ========== 使用示例 ==========
if __name__ == "__main__":
    # 读取点云
    pcd = o3d.io.read_point_cloud(r"E:\point_distance\data\part_car_head.ply")

    # 直通滤波（Z轴）
    points = np.asarray(pcd.points)
    mask = (points[:, 2] >= 0.0) & (points[:, 2] <= 2.0)
    filtered_points = points[mask]

    filtered_pcd = o3d.geometry.PointCloud()
    filtered_pcd.points = o3d.utility.Vector3dVector(filtered_points)

    # 计算指标
    keep_ratio, density_change = compute_keep_ratio_and_density(pcd, filtered_pcd)

    print("=" * 40)
    print("直通滤波指标")
    print("=" * 40)
    print(f"保留率: {keep_ratio:.2f}%")
    print(f"密度变化: {density_change:+.2f}%")

    # 解释
    if density_change > 0:
        print(f"  → 点云变稀疏了 {density_change:.2f}%")
    elif density_change < 0:
        print(f"  → 点云变密集了 {abs(density_change):.2f}%")
    else:
        print(f"  → 密度基本不变")