import numpy as np
import open3d as o3d


def get_ground_truth_seams():
    """返回你手动测量的三维坐标（物体坐标系）。"""
    ground_truth = {
        'up_line1': {
            'start_3d': np.array([142.2, 0, 7.3]),
            'end_3d': np.array([153.9, 0, 149.8])
        },
        'up_line2': {
            'start_3d': np.array([142.2, 0, 7.3]),
            'end_3d': np.array([142.2, 50.3, 7.3])
        },
        'bottom_line1': {
            'start_3d': np.array([8.9, 0, 7.3]),
            'end_3d': np.array([140.2, 0, 7.3])
        },
        'bottom_line2': {
            'start_3d': np.array([142.2, 0, 7.3]),
            'end_3d': np.array([142.2, 50.3, 7.3])
        }
    }
    return ground_truth


def align_and_stitch_seams(reconstructed_seams):
    """
    使用ICP算法将重建的点云对齐到地面真实坐标系，并进行拼接。

    Args:
        reconstructed_seams (dict): 在相机坐标系下重建出的焊缝端点。

    Returns:
        dict: 在物体坐标系下对齐和拼接后的焊缝端点。
    """
    print("\n--- Aligning and Stitching Seams to Ground Truth ---")

    ground_truth = get_ground_truth_seams()

    # --- 1. 对齐上半部分 (up) ---
    # 源点云：重建出的 up_line2 (相机坐标系)
    source_points_up = np.array([
        reconstructed_seams['up_line2']['start_3d'],
        reconstructed_seams['up_line2']['end_3d']
    ])
    source_pcd_up = o3d.geometry.PointCloud()
    source_pcd_up.points = o3d.utility.Vector3dVector(source_points_up)

    # 目标点云：测量的 up_line2 (物体坐标系)
    target_points_up = np.array([
        ground_truth['up_line2']['start_3d'],
        ground_truth['up_line2']['end_3d']
    ])
    target_pcd_up = o3d.geometry.PointCloud()
    target_pcd_up.points = o3d.utility.Vector3dVector(target_points_up)

    print("Aligning 'up' part...")
    # 使用点对点ICP计算变换矩阵 M_up
    # 由于只有两个点，我们可以直接计算一个精确的变换，但用ICP更通用
    # estimate_rigid_transformation 需要点是 (3, N) 的格式
    trans_up = o3d.pipelines.registration.TransformationEstimationPointToPoint().compute_transformation(
        source_pcd_up, target_pcd_up, o3d.utility.Vector2iVector([[0, 0], [1, 1]]))

    print("Transformation matrix for 'up' part (Camera -> Object):")
    print(trans_up)

    # --- 2. 对齐下半部分 (bottom) ---
    # 源点云：重建出的 bottom_line2 (相机坐标系)
    source_points_bottom = np.array([
        reconstructed_seams['bottom_line2']['start_3d'],
        reconstructed_seams['bottom_line2']['end_3d']
    ])
    source_pcd_bottom = o3d.geometry.PointCloud()
    source_pcd_bottom.points = o3d.utility.Vector3dVector(source_points_bottom)

    # 目标点云：测量的 bottom_line2 (物体坐标系)
    target_points_bottom = np.array([
        ground_truth['bottom_line2']['start_3d'],
        ground_truth['bottom_line2']['end_3d']
    ])
    target_pcd_bottom = o3d.geometry.PointCloud()
    target_pcd_bottom.points = o3d.utility.Vector3dVector(target_points_bottom)

    print("\nAligning 'bottom' part...")
    trans_bottom = o3d.pipelines.registration.TransformationEstimationPointToPoint().compute_transformation(
        source_pcd_bottom, target_pcd_bottom, o3d.utility.Vector2iVector([[0, 0], [1, 1]]))

    print("Transformation matrix for 'bottom' part (Camera -> Object):")
    print(trans_bottom)

    # --- 3. 应用变换并组合最终结果 ---
    aligned_seams = {}
    for name, points in reconstructed_seams.items():
        # 创建齐次坐标 (x, y, z, 1)
        start_hom = np.append(points['start_3d'], 1)
        end_hom = np.append(points['end_3d'], 1)

        # 根据焊缝属于 'up' 还是 'bottom' 选择对应的变换矩阵
        if 'up' in name:
            transformed_start = (trans_up @ start_hom.T)[:3]
            transformed_end = (trans_up @ end_hom.T)[:3]
        elif 'bottom' in name:
            transformed_start = (trans_bottom @ start_hom.T)[:3]
            transformed_end = (trans_bottom @ end_hom.T)[:3]
        else:
            continue

        aligned_seams[name] = {
            'start_3d': transformed_start,
            'end_3d': transformed_end
        }

    return aligned_seams, ground_truth