import itertools

import cv2
import numpy as np
from .reconstruction import get_camera_parameters  # 我们仍然需要内参


def get_ground_truth_seams():
    """
    【V5 - 基于图片和新坐标的最终版】
    以公共交点为原点 (0,0,0) 建立坐标系。
    Y轴: 沿着中间公共焊缝。
    X轴: 沿着左下焊缝。
    Z轴: 垂直于XY平面向上。
    """
    print("--- INFO: Using new ground truth based on visual inspection. ---")

    # 1. 定义关键点
    p_origin = np.array([0.0, 0.0, 0.0])  # 公共交点，坐标系原点
    p_middle_end = np.array([0.0, 50.3, 0.0])  # 中间焊缝的终点
    p_bottom_start = np.array([-142.2, 0.0, 0.0])  # 左下焊缝的起点 (沿X负半轴)

    # 对于 up_line1，我们需要一个合理的3D坐标。
    # 它从 p_middle_end (0, 50.3, 0) 开始。
    # 假设它主要在Z方向延伸，我们给它一个长度，比如150。
    # 你给的(-11.7, 142.5)可能存在测量误差或坐标系定义偏差。
    # 我们先用一个理想化的、非退化的点来保证算法能工作。
    p_top_end = np.array([0.0, 50.3, 150.0])  # 假设它竖直向上

    ground_truth = {
        # 上半部分拍摄的两条焊缝
        'up_line1': {
            'start_3d': p_middle_end,  # (0, 50.3, 0)
            'end_3d': p_top_end  # (0, 50.3, 150)
        },
        'up_line2': {
            'start_3d': p_origin,  # (0, 0, 0)
            'end_3d': p_middle_end  # (0, 50.3, 0)
        },
        # 下半部分拍摄的两条焊缝
        'bottom_line1': {
            'start_3d': p_bottom_start,  # (-142.2, 0, 0)
            'end_3d': p_origin  # (0, 0, 0)
        },
        'bottom_line2': {  # 与 up_line2 完全相同
            'start_3d': p_origin,
            'end_3d': p_middle_end
        }
    }
    return ground_truth

# 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 estimate_camera_pose(image_points_2d, object_points_3d, camera_side='L'):
    """
    使用 solvePnP 估计相机位姿。

    Args:
        image_points_2d (np.ndarray): 图像上的2D点 (N, 2)。
        object_points_3d (np.ndarray): 对应的物体坐标系下的3D点 (N, 3)。
        camera_side (str): 'L' or 'R', 用于选择相机内参。

    Returns:
        tuple: (rotation_vector, translation_vector) 相机的位姿。
               这是从物体坐标系到相机坐标系的变换。
    """
    cam_L, cam_R, _ = get_camera_parameters()

    if camera_side == 'L':
        camera_matrix = cam_L['K']
        dist_coeffs = cam_L['kc']
    else:
        camera_matrix = cam_R['K']
        dist_coeffs = cam_R['kc']

    # solvePnP 需要 float64 类型的输入
    object_points_3d = np.array(object_points_3d, dtype=np.float64)
    image_points_2d = np.array(image_points_2d, dtype=np.float64)

    # 使用 solvePnP 求解位姿
    # success: 是否成功
    # rvec: 旋转向量 (Rodrigues vector)
    # tvec: 平移向量
    success, rvec, tvec = cv2.solvePnP(object_points_3d, image_points_2d, camera_matrix, dist_coeffs)

    if not success:
        print("Warning: solvePnP failed to estimate camera pose.")
        return None, None

    return rvec, tvec


def reproject_to_object_coords(endpoints_2d_L, endpoints_2d_R, part_type='up'):
    """
    全新的重建流程（V2 - 修正版）：
    1. 确定2D点和3D点之间最可能的对应关系。
    2. 使用正确的对应关系估计左右相机位姿。
    3. 利用双目信息对所有点进行三角化。
    """
    ground_truth = get_ground_truth_seams()
    cam_L_params, cam_R_params, _ = get_camera_parameters()

    # --- 准备 solvePnP 的原始输入数据 ---
    # object_points_3d: 3D真值点列表，顺序固定
    # image_points_2d_L: 识别出的2D点，顺序可能需要调整
    object_points_3d = []
    image_points_2d_L = []

    seam_keys = []  # 记录焊缝的key，方便后续整理

    for line_name in ['line1', 'line2']:
        gt_key = f"{part_type}_{line_name}"
        if gt_key in ground_truth:
            # 添加3D真值点
            object_points_3d.append(ground_truth[gt_key]['start_3d'])
            object_points_3d.append(ground_truth[gt_key]['end_3d'])

            # 添加对应的2D识别点
            key_L = f"{part_type}_l_{line_name}"
            image_points_2d_L.append(endpoints_2d_L[key_L]['start'])
            image_points_2d_L.append(endpoints_2d_L[key_L]['end'])
            seam_keys.append(gt_key)

    # --- 1. 寻找最佳的2D-3D点对应关系 ---
    # 对于每条焊缝（2个点），有2种可能的匹配（正序或反序）
    # 如果有N条焊缝，就有 2^N 种组合
    # 我们有两条焊缝，所以有 2^2 = 4 种组合

    best_reprojection_error = float('inf')
    best_image_points_L = None

    # a, b 分别代表line1和line2是否需要翻转 (0=不翻转, 1=翻转)
    for a, b in itertools.product([0, 1], repeat=2):
        swaps = (a, b)
        current_image_points_L = list(image_points_2d_L)  # 创建一个副本

        # 根据组合，翻转对应焊缝的起点和终点
        if swaps[0]:  # 翻转 line1
            current_image_points_L[0], current_image_points_L[1] = current_image_points_L[1], current_image_points_L[0]
        if swaps[1]:  # 翻转 line2
            current_image_points_L[2], current_image_points_L[3] = current_image_points_L[3], current_image_points_L[2]

        # 使用当前的对应关系，尝试估计位姿
        rvec_L_try, tvec_L_try = estimate_camera_pose(current_image_points_L, object_points_3d, 'L')

        if rvec_L_try is not None:
            # 计算重投影误差来评估这个组合的好坏
            projected_points, _ = cv2.projectPoints(np.array(object_points_3d), rvec_L_try, tvec_L_try,
                                                    cam_L_params['K'], cam_L_params['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_image_points_L = current_image_points_L

    if best_image_points_L is None:
        print(f"Error: Could not find a valid pose for '{part_type}' part.")
        return None

    print(f"Found best point correspondence for '{part_type}' with reprojection error: {best_reprojection_error:.2f}")

    # --- 2. 使用最佳对应关系，重新进行完整的重建流程 ---

    # 纠正右相机2D点的顺序
    # 这一步有点复杂，我们先假设左右相机的起点/终点翻转是一致的
    # 这是个合理的假设，因为相机离得很近，看到的几何方向应该一样
    best_image_points_R = []
    for line_name in ['line1', 'line2']:
        key_R = f"{part_type}_r_{line_name}"
        points = [endpoints_2d_R[key_R]['start'], endpoints_2d_R[key_R]['end']]
        # 检查左相机的点是否被翻转了
        original_L = [endpoints_2d_L[key_R.replace('_r_', '_l_')]['start'],
                      endpoints_2d_L[key_R.replace('_r_', '_l_')]['end']]
        idx = 0 if line_name == 'line1' else 2
        # 如果左边翻转了，右边也翻转
        if best_image_points_L[idx] != original_L[0]:
            points.reverse()
        best_image_points_R.extend(points)

    # 估计左相机位姿
    rvec_L, tvec_L = estimate_camera_pose(best_image_points_L, object_points_3d, 'L')
    R_L, _ = cv2.Rodrigues(rvec_L)
    P_L = cam_L_params['K'] @ np.hstack((R_L, tvec_L))

    # 估计右相机位姿
    rvec_R, tvec_R = estimate_camera_pose(best_image_points_R, object_points_3d, 'R')
    R_R, _ = cv2.Rodrigues(rvec_R)
    P_R = cam_R_params['K'] @ np.hstack((R_R, tvec_R))

    # 三角化
    points_2d_L_undistorted = cv2.undistortPoints(np.array(best_image_points_L, dtype=np.float32), cam_L_params['K'],
                                                  cam_L_params['kc'], P=cam_L_params['K'])
    points_2d_R_undistorted = cv2.undistortPoints(np.array(best_image_points_R, dtype=np.float32), cam_R_params['K'],
                                                  cam_R_params['kc'], P=cam_R_params['K'])

    points_4d = cv2.triangulatePoints(P_L, P_R, points_2d_L_undistorted.reshape(-1, 2).T,
                                      points_2d_R_undistorted.reshape(-1, 2).T)
    points_3d_object = (points_4d[:3] / points_4d[3]).T

    # 整理输出
    final_seams = {}
    for i, key in enumerate(seam_keys):
        final_seams[key] = {
            'start_3d': points_3d_object[i * 2],
            'end_3d': points_3d_object[i * 2 + 1]
        }

    return final_seams