import numpy as np
import cv2
import open3d as o3d

def get_camera_parameters():
    """
    存储并返回你师兄提供的相机标定参数。
    将所有列表转换为Numpy数组，方便后续计算。
    """
    # 左相机内参
    cam_params_L = {
        'fc': np.array([3774.896, 3770.590]),
        'cc': np.array([1327.950, 956.597]),
        'kc': np.array([-0.098, 0.208, -0.00005, 0.00111, 0]),
        # OpenCV相机矩阵格式 [fx, 0, cx; 0, fy, cy; 0, 0, 1]
        'K': np.array([
            [3774.896, 0, 1327.950],
            [0, 3770.590, 956.597],
            [0, 0, 1]
        ])
    }

    # 右相机内参
    cam_params_R = {
        'fc': np.array([3758.657, 3763.935]),
        'cc': np.array([1274.940, 965.722]),
        'kc': np.array([0.093, -0.219, 0.00079, 0.00255, 0]),
        'K': np.array([
            [3758.657, 0, 1274.940],
            [0, 3763.935, 965.722],
            [0, 0, 1]
        ])
    }

    # 外参 (右相机相对于左相机的变换)
    extrinsics = {
        'R': np.array([
            [0.1169, 0.6292, 0.7683],
            [0.9881, 0.0036, 0.1534],
            [0.0993, -0.7771, -0.6214]
        ]),
        'T': np.array([-220.36786, 2.23290, 30.06279]).reshape(3, 1)  # 平移向量
    }

    return cam_params_L, cam_params_R, extrinsics


def reconstruct_points(points_L, points_R, image_size=(4000, 3000)):
    """
    使用OpenCV进行三维重建的核心函数。

    Args:
        points_L (list of tuples): 左相机图像上的2D点 [(u1, v1), (u2, v2), ...]。
        points_R (list of tuples): 右相机图像上对应的2D点 [(u1, v1), (u2, v2), ...]。
        image_size (tuple): 原始图像的尺寸 (宽度, 高度)，用于立体校正。

    Returns:
        np.ndarray: 重建出的三维点坐标 (N, 3)，单位与标定时使用的单位一致（通常是mm）。
    """
    # 1. 获取相机参数
    cam_L, cam_R, extrinsics = get_camera_parameters()

    # 2. 对输入的2D点进行去畸变
    # 注意：cv2.undistortPoints 需要的格式是 (N, 1, 2) 且为 float32
    points_L_np = np.array(points_L, dtype=np.float32).reshape(-1, 1, 2)
    points_R_np = np.array(points_R, dtype=np.float32).reshape(-1, 1, 2)

    points_L_undistorted = cv2.undistortPoints(points_L_np, cam_L['K'], cam_L['kc'], P=cam_L['K'])
    points_R_undistorted = cv2.undistortPoints(points_R_np, cam_R['K'], cam_R['kc'], P=cam_R['K'])

    # 3. 计算立体校正的投影矩阵
    # stereoRectify 返回很多矩阵，我们只需要P1和P2（新的投影矩阵）
    # 这里我们不需要对图像进行remap，因为我们只关心几个点的变换
    # 注意：这里的R和T是右相机到左相机的变换，与OpenCV的定义一致
    R1, R2, P1, P2, Q, _, _ = cv2.stereoRectify(
        cameraMatrix1=cam_L['K'],
        distCoeffs1=cam_L['kc'],
        cameraMatrix2=cam_R['K'],
        distCoeffs2=cam_R['kc'],
        imageSize=image_size,
        R=extrinsics['R'],
        T=extrinsics['T'].flatten()  # T需要是1D数组
    )

    # 4. 使用 triangulatePoints 进行三角化测量
    # 这个函数需要去畸变后的点和新的投影矩阵
    # 输入点格式需要是 (2, N)
    points_L_for_triangulate = points_L_undistorted.reshape(-1, 2).T
    points_R_for_triangulate = points_R_undistorted.reshape(-1, 2).T

    # triangulatePoints 返回齐次坐标 (4, N)
    points_4d_hom = cv2.triangulatePoints(P1, P2, points_L_for_triangulate, points_R_for_triangulate)

    # 5. 将齐次坐标转换为非齐次坐标
    # 通过除以第四个分量 w
    points_3d = points_4d_hom[:3, :] / points_4d_hom[3, :]

    # 返回转置后的结果，形状为 (N, 3)
    return points_3d.T


def visualize_reconstructed_seams(reconstructed_seams_3d):
    """
    使用 Open3D 可视化重建出的三维焊缝线段。

    Args:
        reconstructed_seams_3d (dict): 包含三维端点坐标的字典。
    """
    print("\n--- Visualizing Final 3-Seam Model vs. Ground Truth ---")

    # 最终的颜色映射
    color_map = {
        # 最终模型 (亮色)
        'bottom_left_final': [1, 0, 0],  # 红色
        'middle_final': [0, 1, 0],  # 绿色
        'top_left_final': [0, 0, 1],  # 蓝色
        # 地面真值 (用稍暗或不同的颜色)
        'bottom_left_truth': [0.8, 0.4, 0.4],  # 粉红
        'middle_truth': [0.4, 0.8, 0.4],  # 浅绿
        'top_left_truth': [0.4, 0.4, 0.8],  # 浅蓝
    }

    geometries = []
    coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=50, origin=[0, 0, 0])
    geometries.append(coordinate_frame)
    # 遍历所有重建出的焊缝
    for name, points in reconstructed_seams_3d.items():
        start_pt = points['start_3d']
        end_pt = points['end_3d']

        # Open3D 需要点和线的列表
        line_points = [start_pt, end_pt]
        line_indices = [[0, 1]]  # 将第一个点和第二个点连接起来
        line_color = color_map.get(name, [0.5, 0.5, 0.5])  # 如果没有定义颜色，则为灰色

        # 创建LineSet对象
        line_set = o3d.geometry.LineSet(
            points=o3d.utility.Vector3dVector(line_points),
            lines=o3d.utility.Vector2iVector(line_indices)
        )
        # 为该线段设置颜色
        line_set.colors = o3d.utility.Vector3dVector([line_color])

        geometries.append(line_set)

        # (可选) 在端点处创建小球体以突出显示
        start_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=10)  # 半径可以调整
        start_sphere.translate(start_pt)
        start_sphere.paint_uniform_color(line_color)
        geometries.append(start_sphere)

        end_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=10)
        end_sphere.translate(end_pt)
        end_sphere.paint_uniform_color(line_color)
        geometries.append(end_sphere)

    # 绘制所有几何对象
    o3d.visualization.draw_geometries(
        geometries,
        window_name="Reconstructed 3D Weld Seams",
        width=1280,
        height=720
    )