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


def evaluate_voxel_filter(pcd, leaf_size=0.05):
    """
    评估体素滤波效果，计算各项指标

    Args:
        pcd: 原始点云
        leaf_size: 体素边长

    Returns:
        metrics: 指标字典
        filtered_pcd: 滤波后点云
    """
    # 执行体素滤波
    filtered_pcd = pcd.voxel_down_sample(leaf_size)

    n_orig = len(pcd.points)
    n_filt = len(filtered_pcd.points)

    # 1. 基础统计指标
    keep_ratio = n_filt / n_orig * 100
    downsample_rate = (n_orig - n_filt) / n_orig * 100
    compression_ratio = n_orig / n_filt if n_filt > 0 else 0

    # 2. 密度变化指标
    points_orig = np.asarray(pcd.points)
    points_filt = np.asarray(filtered_pcd.points)

    # 计算平均最近邻距离（采样加速）
    sample_size = min(1000, n_filt)
    if n_filt > sample_size:
        idx = np.random.choice(n_filt, sample_size, replace=False)
        points_sample = points_filt[idx]
    else:
        points_sample = points_filt

    if len(points_sample) > 5:
        tree = KDTree(points_sample)
        dist, _ = tree.query(points_sample, k=2)
        avg_distance = np.mean(dist[:, 1])
    else:
        avg_distance = 0

    # 3. 输出指标
    metrics = {
        'leaf_size': leaf_size,
        'n_orig': n_orig,
        'n_filt': n_filt,
        'keep_ratio': keep_ratio,
        'downsample_rate': downsample_rate,
        'compression_ratio': compression_ratio,
        'avg_point_distance': avg_distance,
        'estimated_density': 1 / (avg_distance ** 3) if avg_distance > 0 else 0
    }

    return metrics, filtered_pcd


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

    # 测试不同体素大小
    leaf_sizes = [2, 15, 20, 2]

    print("=" * 60)
    print("体素滤波指标评估")
    print("=" * 60)
    print(f"{'体素边长':<10} {'保留率':<12} {'下采样率':<12} {'压缩比':<10} {'平均间距':<12}")
    print("-" * 60)

    for leaf_size in leaf_sizes:
        metrics, filtered = evaluate_voxel_filter(pcd, leaf_size)
        print(f"{leaf_size:<10} {metrics['keep_ratio']:>8.2f}%    "
              f"{metrics['downsample_rate']:>8.2f}%    "
              f"{metrics['compression_ratio']:>8.1f}x    "
              f"{metrics['avg_point_distance']:>8.4f}m")

    # 可视化对比
    print("\n可视化对比...")
    orig_colored = o3d.geometry.PointCloud()
    orig_colored.points = pcd.points
    orig_colored.paint_uniform_color([0, 0, 1])  # 蓝色
    # o3d.visualization.draw_geometries([orig_colored],window_name=f"orignal (blue,{len(pcd.points)}")
    metrics, filtered = evaluate_voxel_filter(pcd, leaf_size=20)
    filtered.paint_uniform_color([1, 0, 0])  # 红色
    # o3d.visualization.draw_geometries([filtered], window_name=f"Voxel (red,{len(filtered.points)}")
    o3d.visualization.draw_geometries(
        [orig_colored, filtered],
        window_name=f"Original (Blue, {len(pcd.points)} pts) vs Voxel (Red, {len(filtered.points)} pts)"
    )