import numpy as np
from scipy.spatial import KDTree
from sklearn.neighbors import NearestNeighbors


def estimate_normals(points, k=20):
    """用 PCA 估计法线"""
    n_points = len(points)
    normals = np.zeros_like(points)

    nbrs = NearestNeighbors(n_neighbors=k).fit(points)
    distances, indices = nbrs.kneighbors(points)

    for i in range(n_points):
        neighbors = points[indices[i]]
        centroid = np.mean(neighbors, axis=0)
        cov = np.cov((neighbors - centroid).T)
        eig_vals, eig_vecs = np.linalg.eig(cov)
        normal = eig_vecs[:, np.argmin(eig_vals)]
        # 统一法线方向
        if normal[2] < 0:
            normal = -normal
        normals[i] = normal

    return normals


def compute_curvature(points, normals, k=20):
    """计算曲率（特征值比值）"""
    curvature = np.zeros(len(points))
    nbrs = NearestNeighbors(n_neighbors=k).fit(points)
    distances, indices = nbrs.kneighbors(points)

    for i in range(len(points)):
        neighbors = points[indices[i]]
        centroid = np.mean(neighbors, axis=0)
        cov = np.cov((neighbors - centroid).T)
        eig_vals, _ = np.linalg.eig(cov)
        curvature[i] = np.min(eig_vals) / (np.sum(eig_vals) + 1e-8)

    return curvature


def region_growing(points, normals, curvature,
                   angle_threshold=30.0,  # 法线夹角阈值（度）
                   curvature_threshold=0.05,  # 曲率阈值
                   min_cluster_size=100):
    """
    区域生长算法
    Args:
        points: (N, 3) 点云坐标
        normals: (N, 3) 法线
        curvature: (N,) 曲率
        angle_threshold: 法线夹角阈值（度）
        curvature_threshold: 曲率阈值
        min_cluster_size: 最小簇大小
    Returns:
        labels: (N,) 每个点的簇标签（-1 表示未分类/噪声）
    """
    n_points = len(points)
    labels = np.full(n_points, -1, dtype=int)

    # 按曲率升序排序（曲率小的作为种子）
    sorted_indices = np.argsort(curvature)

    # 构建 KD-Tree 用于快速邻域搜索
    tree = KDTree(points)

    label_id = 0
    angle_threshold_rad = np.radians(angle_threshold)
    cos_threshold = np.cos(angle_threshold_rad)

    for seed_idx in sorted_indices:
        if labels[seed_idx] != -1:
            continue

        # 开始一个新区域
        current_region = []
        seeds = [seed_idx]
        labels[seed_idx] = label_id

        while seeds:
            current_seed = seeds.pop(0)
            current_region.append(current_seed)

            # 搜索半径内的邻居（动态半径或固定半径）
            neighbors = tree.query_ball_point(points[current_seed], r=0.2)

            for neighbor in neighbors:
                if labels[neighbor] != -1:
                    continue

                # 检查法线夹角
                dot_product = np.abs(np.dot(normals[current_seed], normals[neighbor]))
                if dot_product > cos_threshold:
                    labels[neighbor] = label_id

                    # 如果曲率足够小，作为新种子
                    if curvature[neighbor] < curvature_threshold:
                        seeds.append(neighbor)

        # 检查区域大小
        region_size = np.sum(labels == label_id)
        if region_size < min_cluster_size:
            # 区域太小，标记为噪声
            labels[labels == label_id] = -1
        else:
            label_id += 1

    return labels


# 使用示例
if __name__ == "__main__":
    import open3d as o3d

    # 读取点云
    pcd = o3d.io.read_point_cloud("E:\point_distance\data\\bell_pepper_6.pcd")
    points = np.asarray(pcd.points)

    # 预处理：体素下采样 + 去噪
    pcd = pcd.voxel_down_sample(0.0002)
    pcd, _ = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
    points = np.asarray(pcd.points)

    # 计算法线和曲率
    print("计算法线...")
    normals = estimate_normals(points, k=30)

    print("计算曲率...")
    curvature = compute_curvature(points, normals, k=30)

    # 区域生长分割
    print("区域生长...")
    labels = region_growing(points, normals, curvature,
                            angle_threshold=25.0,
                            curvature_threshold=0.05,
                            min_cluster_size=100)

    print(f"分割完成，共 {len(np.unique(labels)) - (1 if -1 in labels else 0)} 个区域")

    # 可视化
    colors = np.random.rand(len(np.unique(labels)), 3)
    point_colors = np.array([colors[label] if label != -1 else [0.5, 0.5, 0.5]
                             for label in labels])
    pcd.colors = o3d.utility.Vector3dVector(point_colors)
    o3d.visualization.draw_geometries([pcd])