import copy
from scipy.spatial.transform import Rotation as R
import trimesh
import numpy as np
import open3d as o3d
import time
def rotation_matrix_scipy(roll_deg=10, pitch_deg=10, yaw_deg=10):
    """
    使用 scipy 生成绕 XYZ 轴旋转的矩阵
    """
    r = R.from_euler('xyz', [roll_deg, pitch_deg, yaw_deg], degrees=True)
    return r.as_matrix()




def random_rotation_matrix():
    """生成随机旋转矩阵"""
    angle = np.random.uniform(0, 2 * np.pi)
    axis = np.random.randn(3)
    axis = axis / np.linalg.norm(axis)

    K = np.array([
        [0, -axis[2], axis[1]],
        [axis[2], 0, -axis[0]],
        [-axis[1], axis[0], 0]
    ])
    R = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * K @ K
    return R


def random_translation_vector(range_xyz=0.5):
    """生成随机平移向量"""
    return np.random.uniform(-range_xyz, range_xyz, 3)


def transform_point_cloud(pcd, R, t):
    """对点云应用变换"""
    transformed_pcd = copy.deepcopy(pcd)
    points = np.asarray(transformed_pcd.points)
    points_transformed = (R @ points.T).T + t
    transformed_pcd.points = o3d.utility.Vector3dVector(points_transformed)
    return transformed_pcd


def registration_open3d(source_pcd, target_pcd, voxel_size=0.005):
    """
    点云配准（RANSAC + ICP）

    Returns:
        transformation: 4x4 变换矩阵
        fitness: 配准 fitness 分数
        rmse: 配准 RMSE
    """
    # 1. 下采样
    source_down = source_pcd.voxel_down_sample(voxel_size)
    target_down = target_pcd.voxel_down_sample(voxel_size)

    # 2. 估计法线
    radius = voxel_size * 2
    source_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=30))
    target_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=30))

    # 3. 计算 FPFH 特征
    radius_feature = voxel_size * 5
    source_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
        source_down, o3d.geometry.KDTreeSearchParamHybrid(radius=radius_feature, max_nn=100))
    target_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
        target_down, o3d.geometry.KDTreeSearchParamHybrid(radius=radius_feature, max_nn=100))

    # 4. RANSAC 粗配准
    distance_threshold = voxel_size * 1.5
    result_ransac = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
        source_down, target_down, source_fpfh, target_fpfh,
        mutual_filter=True,
        max_correspondence_distance=distance_threshold,
        estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(),
        ransac_n=3,
        checkers=[
            o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
            o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold)
        ],
        criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(100000, 0.999))

    print(f"RANSAC 配准完成")
    source_pcd.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
    target_pcd.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
    # 5. ICP 精配准
    result_icp = o3d.pipelines.registration.registration_icp(
        source_pcd, target_pcd, distance_threshold / 2, result_ransac.transformation,
        o3d.pipelines.registration.TransformationEstimationPointToPlane(),
        o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=200))

    print(f"ICP 精配准完成")
    print(f"  Fitness: {result_icp.fitness:.4f}")
    print(f"  RMSE: {result_icp.inlier_rmse:.6f} m")

    return result_icp.transformation, result_icp.fitness, result_icp.inlier_rmse


def visualize_registration_result(original_pcd, transformed_pcd, registered_pcd, title="Registration Result"):
    """
    可视化原始点云、变换后点云和配准后点云

    Args:
        original_pcd: 原始点云（源）
        transformed_pcd: 应用随机变换后的点云（目标）
        registered_pcd: 配准后的点云
        title: 窗口标题
    """
    # 复制并着色
    source_colored = copy.deepcopy(original_pcd)
    target_colored = copy.deepcopy(transformed_pcd)
    registered_colored = copy.deepcopy(registered_pcd)

    # 设置颜色
    source_colored.paint_uniform_color([0, 0, 1])  # 蓝色 - 原始点云
    target_colored.paint_uniform_color([0, 1, 0])  # 绿色 - 目标点云
    registered_colored.paint_uniform_color([1, 0, 0])  # 红色 - 配准后点云

    # 创建坐标系（可选）
    # coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.5)

    # 可视化
    o3d.visualization.draw_geometries(
        [source_colored, target_colored],
        window_name="Before Registration (Blue: Source, Green: Target)",
        width=1024,
        height=768
    )

    o3d.visualization.draw_geometries(
        [registered_colored, target_colored],
        window_name="After Registration (Red: Registered, Green: Target)",
        width=1024,
        height=768
    )


def visualize_registration_comparison(original_pcd, transformed_pcd, registered_pcd):
    """
    并排可视化配准前后的对比

    Args:
        original_pcd: 原始点云
        transformed_pcd: 变换后的点云（目标）
        registered_pcd: 配准后的点云
    """
    # 复制并着色
    source_before = copy.deepcopy(original_pcd)
    source_before.paint_uniform_color([0, 0, 1])  # 蓝色

    target_before = copy.deepcopy(transformed_pcd)
    target_before.paint_uniform_color([0, 1, 0])  # 绿色

    registered_after = copy.deepcopy(registered_pcd)
    registered_after.paint_uniform_color([1, 0, 0])  # 红色

    target_after = copy.deepcopy(transformed_pcd)
    target_after.paint_uniform_color([0, 1, 0])  # 绿色

    # 创建坐标系
    coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.5)

    # 方法1：分别显示两个窗口
    o3d.visualization.draw_geometries(
        [source_before],
        window_name="Before Registration (Blue: Source, Green: Target)",
        width=1024,
        height=768
    )
    o3d.visualization.draw_geometries(
         [target_before],
        window_name="Before Registration (Blue: Source, Green: Target)",
        width=1024,
        height=768
    )
    o3d.visualization.draw_geometries(
        [registered_after, source_before],
        window_name="After Registration (Red: Registered, Green: Target)",
        width=1024,
        height=768
    )




def compute_point_to_mesh_distance(point_cloud, mesh):
    """
    计算点云到 Mesh 的距离

    Args:
        point_cloud: Open3D 点云对象
        mesh: Open3D 网格对象

    Returns:
        distances: 每个点到最近三角形面的距离
        mean_distance: 平均距离
        rmse: 均方根误差
        max_distance: 最大距离
    """
    # 计算点到 mesh 的距离
    distances = mesh.compute_point_cloud_distance(point_cloud)
    distances = np.asarray(distances)

    mean_dist = np.mean(distances)
    rmse = np.sqrt(np.mean(distances ** 2))
    max_dist = np.max(distances)
    std_dist = np.std(distances)

    return distances, mean_dist, rmse, max_dist, std_dist


# 使用示例
def demo(pcd,mesh):
    # 加载点云和 mesh
    # 计算距离
    distances, mean_dist, rmse, max_dist, std_dist = compute_point_to_mesh_distance(pcd, mesh)

    print(f"平均距离: {mean_dist:.6f} m ({mean_dist * 1000:.3f} mm)")
    print(f"RMSE: {rmse:.6f} m ({rmse * 1000:.3f} mm)")
    print(f"最大距离: {max_dist:.6f} m ({max_dist * 1000:.3f} mm)")
    print(f"标准差: {std_dist:.6f} m ({std_dist * 1000:.3f} mm)")

    # 可视化（按距离着色）
    pcd.colors = o3d.utility.Vector3dVector(get_color_by_distance(distances))
    o3d.visualization.draw_geometries([pcd, mesh])


def get_color_by_distance(distances, max_dist=None):
    """
    根据距离生成颜色（红：远，蓝：近）
    """
    if max_dist is None:
        max_dist = np.percentile(distances, 95)  # 使用95分位数

    colors = np.zeros((len(distances), 3))
    norm_dist = np.clip(distances / (max_dist + 1e-8), 0, 1)

    # 红色通道（距离越远越红）
    colors[:, 0] = norm_dist
    # 蓝色通道（距离越近越蓝）
    colors[:, 2] = 1 - norm_dist

    return colors
# ========== 主程序 ==========
if __name__ == "__main__":
    # 1. 读取原始点云
    file_path = "E:\point_distance\data\\40_noisy.ply"

    source_pcd = o3d.io.read_point_cloud(file_path)

    # points=np.asarray(source_pcd.points)

    t0=time.time()
    if len(source_pcd.points) == 0:
        print("读取失败，请检查文件路径")
        exit()

    print(f"原始点云点数: {len(source_pcd.points)}")

    # 2. 生成随机变换
    # R = random_rotation_matrix()
    R = rotation_matrix_scipy(10, 10, 10)
    t = random_translation_vector(range_xyz=0)

    print(f"随机旋转矩阵:\n{R}")
    print(f"随机平移向量: {t}")

    # 3. 创建目标点云（变换后的源点云）
    target_pcd = transform_point_cloud(source_pcd, R, t)

    # 4. 点云配准
    print("\n开始点云配准...")
    transformation, fitness, rmse = registration_open3d(source_pcd, target_pcd, voxel_size=10)

    # 5. 应用配准变换
    # registered_pcd = copy.deepcopy(source_pcd)
    # registered_pcd.transform(transformation)
    inv_transformation = np.linalg.inv(transformation)
    # 将目标点云变换回源点云坐标系
    target_recovered = copy.deepcopy(target_pcd)
    target_recovered.transform(inv_transformation)


    t1 = time.time()
    print(f"\n配准结果:")
    print(f"  Fitness: {fitness:.4f}")
    print(f"  RMSE: {rmse:.6f} m")
    print(f"  估计的变换矩阵:\n{transformation}")
    print(f"配准time:\n{t1-t0}")


    t2=time.time()
    # 1. 读取数据

    points = np.asarray(target_recovered.points)
    mesh_trimesh = trimesh.load("E:\point_distance\data\\40.stl")
    # ========== 1. 读取数据 ==========
    # 2. 使用 trimesh 计算每个点到 Mesh 的最近距离
    # trimesh 的 nearest 方法非常高效且精确
    closest_points, distances, triangle_ids = mesh_trimesh.nearest.on_surface(points)
    t3 = time.time()

    # distances 即为所求的距离数组
    print(f"平均距离: {np.mean(distances)}")
    print(f"最大距离: {np.max(distances)}")
    print(f"  距离计算:   {t3 - t2:.4f} 秒")
    # 如果需要将结果存回 Open3D 点云进行可视化


    # 6. 可视化对比
    print("\n开始可视化...")
    visualize_registration_comparison(source_pcd, target_pcd,target_recovered)