import numpy as np
from scipy.spatial import KDTree
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


def point_to_line_distance(point, line_start, line_end):
    """
    计算点到三维线段的最短距离

    Args:
        point: 待计算点 (x, y, z)
        line_start: 线段起点 (x, y, z)
        line_end: 线段终点 (x, y, z)

    Returns:
        点到线段的垂直距离
    """
    p = np.array(point)
    a = np.array(line_start)
    b = np.array(line_end)

    # 计算向量
    ab = b - a
    ap = p - a

    # 投影参数 t
    t = np.dot(ap, ab) / np.dot(ab, ab)

    # 找到线段上的最近点
    if t <= 0:
        closest = a
    elif t >= 1:
        closest = b
    else:
        closest = a + t * ab

    # 计算距离
    distance = np.linalg.norm(p - closest)
    return distance


def point_to_plane_distance(point, plane_point1, plane_point2, plane_point3):
    """
    计算点到三维平面的距离（三点确定一个平面）

    Args:
        point: 待计算点 (x, y, z)
        plane_point1, plane_point2, plane_point3: 平面上的三个点

    Returns:
        点到平面的垂直距离
    """
    p = np.array(point)
    p1 = np.array(plane_point1)
    p2 = np.array(plane_point2)
    p3 = np.array(plane_point3)

    # 计算平面法向量
    v1 = p2 - p1
    v2 = p3 - p1
    normal = np.cross(v1, v2)
    normal = normal / np.linalg.norm(normal)

    # 计算点到平面的距离
    distance = abs(np.dot(normal, p - p1))
    return distance


def douglas_peucker_3d(points, epsilon, use_plane=False):
    """
    3D Douglas-Peucker 算法

    Args:
        points: 三维点列表或numpy数组，形状 (N, 3)
        epsilon: 容差阈值
        use_plane: 是否使用平面距离（True）还是线段距离（False）

    Returns:
        简化后的点云列表
    """
    if len(points) <= 2:
        return points.tolist() if isinstance(points, np.ndarray) else points

    # 转换为列表格式
    if isinstance(points, np.ndarray):
        points = points.tolist()

    # 首尾点
    start = points[0]
    end = points[-1]

    # 找到距离最大的点
    max_dist = 0
    max_index = 0

    for i in range(1, len(points) - 1):
        if use_plane and len(points) > 2:
            # 使用平面距离（需要至少3个点确定平面）
            if i > 0 and i < len(points) - 1:
                # 使用当前点、前一个点和后一个点确定平面
                plane_p1 = points[i - 1]
                plane_p2 = points[i]
                plane_p3 = points[i + 1]
                dist = point_to_plane_distance(points[i], plane_p1, plane_p2, plane_p3)
            else:
                dist = point_to_line_distance(points[i], start, end)
        else:
            # 使用线段距离
            dist = point_to_line_distance(points[i], start, end)

        if dist > max_dist:
            max_dist = dist
            max_index = i

    # 递归简化
    if max_dist > epsilon:
        # 保留该点，递归处理左右两部分
        left = douglas_peucker_3d(points[:max_index + 1], epsilon, use_plane)
        right = douglas_peucker_3d(points[max_index:], epsilon, use_plane)

        # 合并结果（避免重复端点）
        return left[:-1] + right
    else:
        # 只保留首尾点
        return [start, end]


class DouglasPeucker3D:
    """
    3D Douglas-Peucker 算法的面向对象封装
    支持点云简化、轨迹压缩、轮廓提取等
    """

    def __init__(self, epsilon=0.01, use_plane=False):
        """
        Args:
            epsilon: 容差阈值
            use_plane: 是否使用平面距离（更适合曲面点云）
        """
        self.epsilon = epsilon
        self.use_plane = use_plane

    def simplify(self, points):
        """
        简化点云

        Args:
            points: (N, 3) 的 numpy 数组或列表

        Returns:
            简化后的点云 (N', 3) numpy 数组
        """
        if isinstance(points, np.ndarray):
            points = points.tolist()

        simplified = douglas_peucker_3d(points, self.epsilon, self.use_plane)
        return np.array(simplified)

    def simplify_trajectory(self, trajectory, time_weight=False):
        """
        专门用于轨迹简化（考虑时间维度）

        Args:
            trajectory: (N, 4) 的数组 [x, y, z, t] 或 [x, y, z]
            time_weight: 是否加入时间权重

        Returns:
            简化后的轨迹
        """
        if time_weight and trajectory.shape[1] == 4:
            # 时间加权：将时间作为第四维参与距离计算
            # 需要自定义距离函数
            pass

        return self.simplify(trajectory[:, :3])

    def extract_contour(self, points, return_indices=False):
        """
        提取点云轮廓（结合法线估计）

        Args:
            points: (N, 3) 点云
            return_indices: 是否返回原始点索引

        Returns:
            轮廓点集
        """
        # 先进行粗略简化
        simplified = self.simplify(points)

        # 这里可以结合法线估计进一步提取边界点
        # 简化实现：返回简化后的点作为轮廓候选

        if return_indices:
            # 查找简化点在原始点云中的最近邻索引
            from scipy.spatial import KDTree
            tree = KDTree(points)
            indices = []
            for p in simplified:
                _, idx = tree.query(p)
                indices.append(idx)
            return simplified, np.array(indices)

        return simplified


# ========== 示例和使用 ==========

def generate_test_curve(n_points=100, noise=0.02):
    """
    生成测试用的三维曲线（螺旋线 + 噪声）
    """
    t = np.linspace(0, 4 * np.pi, n_points)
    x = np.sin(t)
    y = np.cos(t)
    z = t / (2 * np.pi)

    points = np.column_stack([x, y, z])

    # 添加噪声
    if noise > 0:
        points += np.random.randn(n_points, 3) * noise

    return points


def visualize_simplification(original, simplified, title="3D Douglas-Peucker Simplification"):
    """
    可视化原始点云和简化后的点云
    """
    fig = plt.figure(figsize=(12, 6))

    # 原始点云
    ax1 = fig.add_subplot(121, projection='3d')
    ax1.scatter(original[:, 0], original[:, 1], original[:, 2],
                c='blue', s=10, alpha=0.6, label='Original')
    ax1.plot(original[:, 0], original[:, 1], original[:, 2],
             'b-', alpha=0.3)
    ax1.set_title(f'Original: {len(original)} points')
    ax1.legend()

    # 简化后点云
    ax2 = fig.add_subplot(122, projection='3d')
    ax2.scatter(simplified[:, 0], simplified[:, 1], simplified[:, 2],
                c='red', s=30, alpha=0.8, label='Simplified')
    ax2.plot(simplified[:, 0], simplified[:, 1], simplified[:, 2],
             'r-', linewidth=2)
    ax2.set_title(f'Simplified: {len(simplified)} points (ε={epsilon})')
    ax2.legend()

    plt.suptitle(title)
    plt.tight_layout()
    plt.show()


def compare_epsilon(points, epsilons):
    """
    比较不同 epsilon 值的简化效果
    """
    results = {}
    for eps in epsilons:
        dp = DouglasPeucker3D(epsilon=eps)
        simplified = dp.simplify(points)
        results[eps] = {
            'points': simplified,
            'count': len(simplified),
            'compression_ratio': len(simplified) / len(points)
        }
        print(f"ε={eps:.4f}: {len(simplified)} 点 (压缩率: {results[eps]['compression_ratio']:.2%})")

    return results


def demo_point_cloud():
    """
    点云简化演示（使用螺旋线模拟）
    """
    print("=" * 50)
    print("3D Douglas-Peucker 算法演示")
    print("=" * 50)

    # 生成测试数据
    points = generate_test_curve(200, noise=0.01)
    print(f"\n原始点云: {len(points)} 个点")

    # 测试不同 epsilon
    epsilons = [0.01, 0.03, 0.05, 0.1, 0.2]
    print("\n不同容差阈值的简化效果:")
    results = compare_epsilon(points, epsilons)

    # 可视化
    dp = DouglasPeucker3D(epsilon=0.05)
    simplified = dp.simplify(points)
    visualize_simplification(points, simplified,
                             f"3D DP Simplification (ε=0.05, {len(simplified)}/{len(points)} points)")


def demo_trajectory():
    """
    轨迹简化演示（模拟移动轨迹）
    """
    print("\n" + "=" * 50)
    print("3D 轨迹简化演示")
    print("=" * 50)

    # 生成测试轨迹（随机路径）
    np.random.seed(42)
    n_points = 100
    t = np.linspace(0, 10, n_points)
    x = np.sin(t) + np.random.randn(n_points) * 0.02
    y = np.cos(t) + np.random.randn(n_points) * 0.02
    z = t / 5 + np.random.randn(n_points) * 0.01

    trajectory = np.column_stack([x, y, z])
    print(f"原始轨迹: {len(trajectory)} 个点")

    # 简化轨迹
    dp = DouglasPeucker3D(epsilon=0.08)
    simplified = dp.simplify_trajectory(trajectory)
    print(f"简化后轨迹: {len(simplified)} 个点 (压缩率: {len(simplified) / len(trajectory):.2%})")

    # 可视化
    visualize_simplification(trajectory, simplified, "3D Trajectory Simplification")


def demo_point_cloud_from_file(pcd_file=None):
    """
    从文件加载点云并简化
    """
    try:
        import open3d as o3d

        if pcd_file is None:
            # 使用示例：生成随机点云
            print("\n生成随机点云进行演示...")
            points = np.random.randn(500, 3) * 0.5
        else:
            pcd = o3d.io.read_point_cloud(pcd_file)
            points = np.asarray(pcd.points)

        print(f"原始点云: {len(points)} 个点")

        # 简化
        dp = DouglasPeucker3D(epsilon=0.05)

        # 注意：传统 DP 算法需要有序点集，对于无序点云需要先排序或使用其他方法
        # 这里演示对有序轨迹/轮廓的应用
        print("\n注意: DP算法适用于有序点集（如轮廓线、轨迹），")
        print("对于无序点云，建议先进行排序或使用其他简化方法（如体素滤波）")

    except ImportError:
        print("需要安装 open3d 来读取点云文件: pip install open3d")


if __name__ == "__main__":
    # 运行演示
    demo_point_cloud()
    demo_trajectory()

    # 如果需要处理实际点云文件
    # demo_point_cloud_from_file("your_pointcloud.pcd")