import os
import cv2
import numpy as np

# 导入必要的模块
from script.yolo_detector import detect_crop_area
from script.linknet_segmentor import segment_and_find_endpoints
from script.final_reconstruction import merge_seams  # 我们依然需要合并函数
from script.reconstruction import visualize_reconstructed_seams  # 和可视化函数


def reconstruct_with_optimized_params(points_L, points_R, calib_data, image_size=(4000, 3000)):
    """
    使用优化好的参数文件，进行高效的标准双目重建。
    返回在左相机坐标系下的三维点。
    """
    # 从标定数据中加载新的内外参
    K_L = np.array([
        [calib_data['optimized_intrinsics_L'][0], 0, calib_data['optimized_intrinsics_L'][2]],
        [0, calib_data['optimized_intrinsics_L'][1], calib_data['optimized_intrinsics_L'][3]],
        [0, 0, 1]
    ])
    kc_L = calib_data['dist_coeffs_L']

    K_R = np.array([
        [calib_data['optimized_intrinsics_R'][0], 0, calib_data['optimized_intrinsics_R'][2]],
        [0, calib_data['optimized_intrinsics_R'][1], calib_data['optimized_intrinsics_R'][3]],
        [0, 0, 1]
    ])
    kc_R = calib_data['dist_coeffs_R']

    # 使用新的、优化过的外参！
    new_extrinsics = calib_data['new_extrinsics'].item()  # .item() 用于从numpy对象数组中提取字典
    R = new_extrinsics['R']
    t = new_extrinsics['t']

    # 标准的立体校正和三角化流程
    R1, R2, P1, P2, _, _, _ = cv2.stereoRectify(K_L, kc_L, K_R, kc_R, image_size, R, t)

    points_L_undistorted = cv2.undistortPoints(np.array(points_L, dtype=np.float32), K_L, kc_L, P=P1)
    points_R_undistorted = cv2.undistortPoints(np.array(points_R, dtype=np.float32), K_R, kc_R, P=P2)

    points_4d_hom = cv2.triangulatePoints(P1, P2, points_L_undistorted.reshape(-1, 2).T,
                                          points_R_undistorted.reshape(-1, 2).T)
    points_3d_camL = (points_4d_hom[:3] / points_4d_hom[3]).T

    return points_3d_camL


def get_transform_from_pose(pose):
    """从6自由度位姿向量计算 4x4 逆变换矩阵（相机->物体）。"""
    rvec, tvec = pose[:3], pose[3:]
    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 run_deployment_pipeline(calib_data):
    """
    最终部署流程：加载标定文件，快速完成重建。
    """
    print("--- Running Deployment Pipeline with Optimized Parameters ---")

    # 1. 运行2D识别 (这部分和之前一样)
    # 你可以从之前的 main.py 复制 run_full_recognition_pipeline 函数过来
    # 或者我们在这里重新写一个简化版的
    from main import run_full_recognition_pipeline  # 假设之前的main.py还在
    all_2d_endpoints = run_full_recognition_pipeline()
    if not all_2d_endpoints:
        print("2D recognition failed. Exiting.")
        return

    reconstructed_4_seams = {}

    # 2. 分别处理 'up' 和 'bottom' 的重建和变换
    for part_type in ['up', 'bottom']:
        print(f"\nProcessing '{part_type}' part...")

        # a. 收集该部分的所有2D点
        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}")

        # b. 使用优化后的参数进行标准双目重建
        points_camL = reconstruct_with_optimized_params(points_L, points_R, calib_data)

        # c. 获取对应的变换矩阵并应用
        # 注意：我们假设两次拍摄时相机与物体的相对关系固定，
        # 因此理论上 up 和 bottom 的变换矩阵应该是一样的。
        # 我们使用 'up' 拍摄时计算出的位姿作为全局基准。
        global_transform = get_transform_from_pose(calib_data['pose_up_L'])

        points_camL_hom = np.hstack([points_camL, np.ones((points_camL.shape[0], 1))])
        points_object = (global_transform @ points_camL_hom.T).T[:, :3]

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

    # 3. 合并为最终的三线模型
    final_3_seam_model = merge_seams(reconstructed_4_seams)

    # 4. 打印和可视化
    print("\n--- Final 3-Seam Model (Object Coordinate System) ---")
    for name, points in final_3_seam_model.items():
        start_str = np.array2string(points['start_3d'], formatter={'float_kind': lambda x: "%.2f" % x})
        end_str = np.array2string(points['end_3d'], formatter={'float_kind': lambda x: "%.2f" % x})
        print(f"{name}: Start={start_str}, End={end_str}")

    # 可视化...
    from script.pose_estimation import get_ground_truth_seams
    ground_truth_data = get_ground_truth_seams()
    comparison_data = {}
    for name, points in final_3_seam_model.items():
        comparison_data[name + '_final'] = points
    comparison_data['bottom_left_truth'] = ground_truth_data['bottom_line1']
    comparison_data['middle_truth'] = ground_truth_data['up_line2']
    comparison_data['top_left_truth'] = ground_truth_data['up_line1']
    visualize_reconstructed_seams(comparison_data)


if __name__ == '__main__':
    # 定义标定文件路径
    calib_file_path = 'optimized_camera_parameters.npz'

    if not os.path.exists(calib_file_path):
        print(f"Error: Calibration file not found at '{calib_file_path}'")
        print("Please run the main.py with the global optimization first to generate this file.")
    else:
        # 加载标定文件
        print(f"Loading optimized parameters from '{calib_file_path}'...")
        calibration_data = np.load(calib_file_path, allow_pickle=True)

        # 运行部署流程
        run_deployment_pipeline(calibration_data)