import cv2
import numpy as np
import itertools
from .reconstruction import get_camera_parameters
from .pose_estimation import get_ground_truth_seams, estimate_camera_pose


def get_global_transform_from_up_data(all_2d_endpoints):
    """
    只使用 'up' 部分的数据，计算一个全局的、唯一的“相机->物体”变换矩阵。
    """
    print("\n--- Calculating Global Transform Matrix using 'up' data ---")
    ground_truth = get_ground_truth_seams()

    # 1. 准备 'up' 部分的数据
    object_points_3d = []
    image_points_2d_L = []

    for line_name in ['line1', 'line2']:
        gt_key = f"up_{line_name}"
        key_L = f"up_l_{line_name}"
        object_points_3d.extend([ground_truth[gt_key]['start_3d'], ground_truth[gt_key]['end_3d']])
        image_points_2d_L.extend([all_2d_endpoints[key_L]['start'], all_2d_endpoints[key_L]['end']])

    # 2. 寻找最佳点对应关系
    best_reprojection_error = float('inf')
    best_pose = None

    for a, b in itertools.product([0, 1], repeat=2):
        current_image_points_L = list(image_points_2d_L)
        if a: current_image_points_L[0], current_image_points_L[1] = current_image_points_L[1], current_image_points_L[
            0]
        if b: current_image_points_L[2], current_image_points_L[3] = current_image_points_L[3], current_image_points_L[
            2]

        rvec, tvec = estimate_camera_pose(current_image_points_L, object_points_3d, 'L')
        if rvec is not None:
            projected_points, _ = cv2.projectPoints(np.array(object_points_3d), rvec, tvec,
                                                    get_camera_parameters()[0]['K'], get_camera_parameters()[0]['kc'])
            error = cv2.norm(np.array(current_image_points_L, dtype=np.float32),
                             projected_points.reshape(-1, 2).astype(np.float32), cv2.NORM_L2)

            if error < best_reprojection_error:
                best_reprojection_error = error
                best_pose = (rvec, tvec)

    if best_pose is None:
        print("Fatal Error: Could not calculate a valid global transform.")
        return None

    print(f"Global transform calculated with reprojection error: {best_reprojection_error:.2f}")

    # 3. 构建 4x4 变换矩阵 (从相机坐标系到物体坐标系)
    rvec, tvec = best_pose
    R_cam_from_obj, _ = cv2.Rodrigues(rvec)
    R_obj_from_cam = R_cam_from_obj.T
    t_obj_from_cam = -R_obj_from_cam @ tvec

    transform_matrix = np.eye(4)
    transform_matrix[:3, :3] = R_obj_from_cam
    transform_matrix[:3, 3] = t_obj_from_cam.flatten()

    return transform_matrix


def reconstruct_in_camera_coords(points_L, points_R, image_size=(4000, 3000)):
    # ... (这个函数保持不变)
    cam_L, cam_R, extrinsics = get_camera_parameters()
    R1, R2, P1, P2, _, _, _ = cv2.stereoRectify(cam_L['K'], cam_L['kc'], cam_R['K'], cam_R['kc'], image_size,
                                                extrinsics['R'], extrinsics['T'].flatten())
    points_L_undistorted = cv2.undistortPoints(np.array(points_L, dtype=np.float32), cam_L['K'], cam_L['kc'], P=P1)
    points_R_undistorted = cv2.undistortPoints(np.array(points_R, dtype=np.float32), cam_R['K'], cam_R['kc'], P=P2)
    points_4d_hom = cv2.triangulatePoints(P1, P2, points_L_undistorted.reshape(-1, 2).T,
                                          points_R_undistorted.reshape(-1, 2).T)
    return (points_4d_hom[:3] / points_4d_hom[3]).T


def final_reconstruction_pipeline(all_2d_endpoints):
    # 1. 计算唯一的、全局的变换矩阵
    global_transform = get_global_transform_from_up_data(all_2d_endpoints)
    if global_transform is None:
        return None

    reconstructed_4_seams = {}

    for part_type in ['up', 'bottom']:
        # 2. 对每个部分进行标准双目重建
        points_L, points_R, seam_keys = [], [], []
        for line_name in ['line1', 'line2']:
            key_L = f"{part_type}_l_{line_name}"
            key_R = f"{part_type}_r_{line_name}"
            points_L.extend([all_2d_endpoints[key_L]['start'], all_2d_endpoints[key_L]['end']])
            points_R.extend([all_2d_endpoints[key_R]['start'], all_2d_endpoints[key_R]['end']])
            seam_keys.append(f"{part_type}_{line_name}")

        points_camL = reconstruct_in_camera_coords(points_L, points_R)

        # 3. 使用同一个全局矩阵进行变换
        points_camL_hom = np.hstack([points_camL, np.ones((points_camL.shape[0], 1))])
        points_object = (global_transform @ points_camL_hom.T).T[:, :3]

        # 4. 整理结果
        for i, key in enumerate(seam_keys):
            reconstructed_4_seams[key] = {'start_3d': points_object[i * 2], 'end_3d': points_object[i * 2 + 1]}

    return reconstructed_4_seams




def merge_seams(reconstructed_seams_dict):
    """
    将重建出的四条焊缝合并为最终的三条焊缝模型。

    Args:
        reconstructed_seams_dict (dict): 包含 'up_line1', 'up_line2',
                                         'bottom_line1', 'bottom_line2' 的字典。

    Returns:
        dict: 包含 'bottom_left', 'middle', 'top_left' 三条最终焊缝的字典。
    """
    print("\n--- Merging seams into final 3-line model ---")

    if not all(k in reconstructed_seams_dict for k in ['up_line1', 'up_line2', 'bottom_line1', 'bottom_line2']):
        print("Error: Missing reconstructed seams for merging.")
        return None

    # 提取所有需要的端点
    bl1_start = reconstructed_seams_dict['bottom_line1']['start_3d']
    bl1_end = reconstructed_seams_dict['bottom_line1']['end_3d']

    ul2_start = reconstructed_seams_dict['up_line2']['start_3d']
    ul2_end = reconstructed_seams_dict['up_line2']['end_3d']

    bl2_start = reconstructed_seams_dict['bottom_line2']['start_3d']
    bl2_end = reconstructed_seams_dict['bottom_line2']['end_3d']

    ul1_start = reconstructed_seams_dict['up_line1']['start_3d']
    ul1_end = reconstructed_seams_dict['up_line1']['end_3d']

    # --- 定义最终的三条线 ---

    # 1. 左下焊缝 (bottom_left)
    # 直接使用 bottom_line1。为了保证方向一致，我们让它从X值较小的点指向X值较大的点。
    bottom_left_points = sorted([bl1_start, bl1_end], key=lambda p: p[0])

    # 2. 中间焊缝 (middle)
    # 这是 up_line2 和 bottom_line2 的合并。理论上它们应该重合。
    # 我们可以取四个点的平均值来得到更稳健的起点和终点。
    # 公共起点应该是 (bl1_end, ul2_start, bl2_start) 的平均值
    middle_start = np.mean([bl1_end, ul2_start, bl2_start], axis=0)
    # 公共终点应该是 (ul2_end, bl2_end, ul1_start) 的平均值
    middle_end = np.mean([ul2_end, bl2_end, ul1_start], axis=0)

    # 3. 左上焊缝 (top_left)
    # 直接使用 up_line1。
    top_left_points = [ul1_start, ul1_end]  # 保持原始方向

    final_3_seams = {
        'bottom_left': {
            'start_3d': bottom_left_points[0],
            'end_3d': bottom_left_points[1]
        },
        'middle': {
            'start_3d': middle_start,
            'end_3d': middle_end
        },
        'top_left': {
            'start_3d': top_left_points[0],
            'end_3d': top_left_points[1]
        }
    }

    return final_3_seams