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.reconstruction import visualize_reconstructed_seams, reconstruct_points
from script.pose_estimation import get_ground_truth_seams, reproject_to_object_coords  # 我们只需要真值
# 导入我们最终的重建流程
from script.final_reconstruction import final_reconstruction_pipeline, merge_seams
import itertools # 确保导入itertools


from script.global_optimizer import run_global_optimization, merge_seams
from script.pose_estimation import get_ground_truth_seams
from script.reconstruction import visualize_reconstructed_seams

def run_full_recognition_pipeline():
    """
    运行完整的识别流程：YOLO定位 -> LinkNet分割 -> 端点提取。
    """
    # 1. 定义路径
    base_dir = os.path.dirname(os.path.abspath(__file__))
    data_map = {
        'up': {
            'l_img': os.path.join(base_dir, 'data', 'origin', 'up', 'l1.jpeg'),
            'r_img': os.path.join(base_dir, 'data', 'origin', 'up', 'r1.jpeg'),
            'yolo_model': os.path.join(base_dir, 'module', 'yolov8', 'up.pt'),
            'linknet_models': {
                'line1': os.path.join(base_dir, 'module', 'linknet', 'best_linknet_up_model_line1.pth'),
                'line2': os.path.join(base_dir, 'module', 'linknet', 'best_linknet_up_model_line2.pth')
            }
        },
        'bottom': {
            'l_img': os.path.join(base_dir, 'data', 'origin', 'bottom', 'l1.jpeg'),
            'r_img': os.path.join(base_dir, 'data', 'origin', 'bottom', 'r1.jpeg'),
            'yolo_model': os.path.join(base_dir, 'module', 'yolov8', 'bottom.pt'),
            'linknet_models': {
                'line1': os.path.join(base_dir, 'module', 'linknet', 'best_linknet_bottom_model_line1.pth'),
                'line2': os.path.join(base_dir, 'module', 'linknet', 'best_linknet_bottom_model_line2.pth')
            }
        }
    }
    output_dir = os.path.join(base_dir, 'data', 'processed')
    os.makedirs(output_dir, exist_ok=True)

    all_endpoints = {}

    for part, paths in data_map.items():
        print(f"\n--- Processing '{part}' part ---")
        for img_path, side in [(paths['l_img'], 'l'), (paths['r_img'], 'r')]:
            print(f"\n-- Analyzing image: {os.path.basename(img_path)} --")

            crop_box = detect_crop_area(img_path, paths['yolo_model'])
            if not crop_box:
                print(f"Skipping further processing for {os.path.basename(img_path)}.")
                continue

            original_image_vis = cv2.imread(img_path)

            for line_name, linknet_path in paths['linknet_models'].items():
                endpoints = segment_and_find_endpoints(original_image_vis, crop_box, linknet_path)

                if endpoints:
                    start_pt, end_pt = endpoints
                    result_key = f"{part}_{side}_{line_name}"
                    all_endpoints[result_key] = {'start': start_pt, 'end': end_pt}

                    # --- 在可视化图像上绘制结果 (增强版) ---
                    # 1. 绘制端点圆圈
                    cv2.circle(original_image_vis, start_pt, 15, (0, 255, 0), -1)  # 绿色起点
                    cv2.circle(original_image_vis, end_pt, 15, (0, 0, 255), -1)  # 红色终点
                    # 2. 绘制连接线
                    cv2.line(original_image_vis, start_pt, end_pt, (255, 0, 0), 4)

                    # 3. 添加文本标签
                    # 计算线段中点作为文本放置位置
                    mid_point = ((start_pt[0] + end_pt[0]) // 2, (start_pt[1] + end_pt[1]) // 2)
                    # 在中点上方放置文本
                    text_pos = (mid_point[0], mid_point[1] - 20)
                    cv2.putText(original_image_vis,
                                result_key,
                                text_pos,
                                cv2.FONT_HERSHEY_SIMPLEX,
                                2,  # 字体大小
                                (255, 255, 0),  # 字体颜色 (青色)
                                4,  # 字体粗细
                                cv2.LINE_AA)

            # 绘制YOLO框并保存最终的可视化结果
            cv2.rectangle(original_image_vis, (crop_box[0], crop_box[1]), (crop_box[2], crop_box[3]), (0, 255, 255), 4)
            save_path = os.path.join(output_dir, f'{part}_{side}_final_result.jpg')
            cv2.imwrite(save_path, original_image_vis)
            print(f"Saved final visualization to {save_path}")
    # 3. 打印总结
    print("\n--- Final Endpoints Summary (in original image coordinates) ---")
    for name, points in all_endpoints.items():
        print(f"{name}: Start={points['start']}, End={points['end']}")

    return all_endpoints


def run_3d_reconstruction(all_2d_endpoints):
    """
    根据识别出的2D端点，重建出三维焊缝。
    """
    print("\n--- Starting 3D Reconstruction ---")

    # 这个字典将存储最终的三维坐标
    reconstructed_seams_3d = {}

    # 需要重建的焊缝对
    # 例如：'up_line1' 对应 up_l_line1 和 up_r_line1
    seam_pairs = ['up_line1', 'up_line2', 'bottom_line1', 'bottom_line2']

    for seam_name in seam_pairs:
        key_L = f"{seam_name.split('_')[0]}_l_{seam_name.split('_')[1]}"  # e.g., 'up_l_line1'
        key_R = f"{seam_name.split('_')[0]}_r_{seam_name.split('_')[1]}"  # e.g., 'up_r_line1'

        # 检查左右相机的点是否都已识别
        if key_L not in all_2d_endpoints or key_R not in all_2d_endpoints:
            print(f"Warning: Missing points for seam '{seam_name}'. Cannot reconstruct.")
            continue

        # 准备输入点列表：[start_point, end_point]
        points_L = [all_2d_endpoints[key_L]['start'], all_2d_endpoints[key_L]['end']]
        points_R = [all_2d_endpoints[key_R]['start'], all_2d_endpoints[key_R]['end']]

        # 调用重建函数
        # 假设你的图像尺寸是 4000x3000，如果不是，请修改
        # 这是一个重要的参数，需要与标定时使用的图像尺寸一致！
        points_3d = reconstruct_points(points_L, points_R, image_size=(4000, 3000))

        reconstructed_seams_3d[seam_name] = {
            'start_3d': points_3d[0],
            'end_3d': points_3d[1]
        }

    # --- 打印最终的三维坐标结果 ---
    print("\n--- Final 3D Seam Endpoints (in Left Camera Coordinate System, unit: mm) ---")
    for name, points in reconstructed_seams_3d.items():
        start_str = np.array2string(points['start_3d'], formatter={'float_kind': lambda x: "%.3f" % x})
        end_str = np.array2string(points['end_3d'], formatter={'float_kind': lambda x: "%.3f" % x})
        print(f"{name}:")
        print(f"  Start 3D: {start_str}")
        print(f"  End 3D:   {end_str}")

    return reconstructed_seams_3d


def run_new_reconstruction_pipeline(all_2d_endpoints):
    """
    使用 solvePnP 的全新重建和拼接流程。
    """
    print("\n--- Starting NEW Reconstruction Pipeline (with solvePnP) ---")

    # --- 处理上半部分 ---
    print("\nProcessing 'up' part...")
    reconstructed_up = reproject_to_object_coords(all_2d_endpoints, all_2d_endpoints, part_type='up')

    # --- 处理下半部分 ---
    print("\nProcessing 'bottom' part...")
    reconstructed_bottom = reproject_to_object_coords(all_2d_endpoints, all_2d_endpoints, part_type='bottom')

    # --- 合并结果 ---
    final_reconstructed_seams = {}
    if reconstructed_up:
        final_reconstructed_seams.update(reconstructed_up)
    if reconstructed_bottom:
        final_reconstructed_seams.update(reconstructed_bottom)

    return final_reconstructed_seams


if __name__ == '__main__':
    final_2d_endpoints = run_full_recognition_pipeline()
    ground_truth = get_ground_truth_seams()

    final_4_seams = {}
    if final_2d_endpoints:
        # 直接调用全局优化
        final_4_seams = run_global_optimization(final_2d_endpoints, ground_truth)

    final_3_seam_model = {}
    if final_4_seams:
        final_3_seam_model = merge_seams(final_4_seams)