#!/usr/bin/env python3
"""
test.py
摄像头畸变校正效果测试 - 仅显示校正前后对比画面
"""

import cv2
import numpy as np
import argparse
import os

def is_rtsp_url(input_str):
    """检查输入是否为RTSP URL"""
    return input_str.startswith('rtsp://')

def load_calibration_params(yaml_file):
    """从YAML文件加载标定参数"""
    if not os.path.exists(yaml_file):
        raise FileNotFoundError(f"标定文件 {yaml_file} 不存在")
    
    fs = cv2.FileStorage(yaml_file, cv2.FILE_STORAGE_READ)
    if not fs.isOpened():
        raise IOError(f"无法打开标定文件 {yaml_file}")
    
    # 读取参数
    camera_matrix = fs.getNode("camera_matrix").mat()
    dist_coeffs = fs.getNode("dist_coeffs").mat()
    
    fs.release()
    
    print(f"✓ 成功加载标定参数")
    print(f"  相机矩阵:\n{camera_matrix}")
    print(f"  畸变系数: {dist_coeffs.flatten()}")
    
    return camera_matrix, dist_coeffs

def main():
    parser = argparse.ArgumentParser(description='摄像头畸变校正效果测试')
    parser.add_argument("-i", "--input", type=str, required=True,
                        help="输入源（摄像头设备索引、视频文件路径或RTSP URL）")
    parser.add_argument("-c", "--calibration", type=str, required=True,
                        help="标定参数YAML文件路径")
    parser.add_argument("-f", "--fisheye", action="store_true",
                        help="是否为鱼眼相机")
    
    args = parser.parse_args()
    
    # 加载标定参数
    try:
        camera_matrix, dist_coeffs = load_calibration_params(args.calibration)
    except Exception as e:
        print(f"❌ 加载标定参数失败: {e}")
        return
    
    # 初始化视频源
    if is_rtsp_url(args.input):
        # RTSP流
        cap = cv2.VideoCapture(args.input, cv2.CAP_FFMPEG)
        try:
            cap.set(cv2.CAP_PROP_FFMPEG_OPTION, "rtsp_transport", "tcp")
            cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
        except:
            pass
    elif args.input.isdigit():
        # 摄像头设备索引
        cap = cv2.VideoCapture(int(args.input))
    else:
        # 视频文件
        cap = cv2.VideoCapture(args.input)
    
    if not cap.isOpened():
        print(f"❌ 无法打开视频源: {args.input}")
        return
    
    # 获取实际分辨率
    actual_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    actual_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    print(f"视频源分辨率: {actual_w}x{actual_h}")
    
    print("✓ 准备开始校正")
    print("按 'q' 键退出")
    print("按 's' 键切换校正方法")
    
    window_name = "畸变校正对比 - 原始(左) vs 校正后(右)"
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
    
    # 校正方法选择
    use_remap = True  # 默认使用remap方法
    
    frame_count = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            print("读取帧失败")
            # 对于RTSP流，尝试重新连接
            if is_rtsp_url(args.input):
                cap.release()
                cap = cv2.VideoCapture(args.input, cv2.CAP_FFMPEG)
                continue
            break
        
        frame_count += 1
        
        # 应用畸变校正
        try:
            if args.fisheye:
                if use_remap:
                    # 使用映射方法
                    new_camera_matrix = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(
                        camera_matrix, dist_coeffs, (actual_w, actual_h), np.eye(3), balance=0.8)
                    map1, map2 = cv2.fisheye.initUndistortRectifyMap(
                        camera_matrix, dist_coeffs, np.eye(3), new_camera_matrix, 
                        (actual_w, actual_h), cv2.CV_16SC2)
                    undistorted = cv2.remap(frame, map1, map2, 
                                          interpolation=cv2.INTER_LINEAR,
                                          borderMode=cv2.BORDER_CONSTANT)
                else:
                    # 使用直接校正方法
                    undistorted = cv2.fisheye.undistortImage(frame, camera_matrix, dist_coeffs)
            else:
                # 普通相机校正
                undistorted = cv2.undistort(frame, camera_matrix, dist_coeffs)
                
        except Exception as e:
            print(f"校正失败: {e}")
            undistorted = frame.copy()  # 失败时使用原图
        
        # 检查校正后的图像是否有效
        if undistorted is None or undistorted.size == 0:
            print(f"警告: 第 {frame_count} 帧校正失败")
            undistorted = frame.copy()
        
        # 调整大小以便显示
        max_height = 600
        scale_orig = max_height / frame.shape[0]
        orig_resized = cv2.resize(frame, 
                                (int(frame.shape[1] * scale_orig), 
                                 int(frame.shape[0] * scale_orig)))
        
        scale_undist = max_height / undistorted.shape[0]
        undistorted_resized = cv2.resize(undistorted, 
                                       (int(undistorted.shape[1] * scale_undist), 
                                        int(undistorted.shape[0] * scale_undist)))
        
        # 确保两个图像高度相同
        min_height = min(orig_resized.shape[0], undistorted_resized.shape[0])
        orig_resized = orig_resized[:min_height, :]
        undistorted_resized = undistorted_resized[:min_height, :]
        
        # 水平拼接显示对比
        comparison = np.hstack((orig_resized, undistorted_resized))
        
        # 添加文本说明
        cv2.putText(comparison, "原始图像 (有畸变)", 
                   (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        cv2.putText(comparison, "校正后图像", 
                   (orig_resized.shape[1] + 10, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        
        # 显示相机类型和校正方法
        camera_type = "鱼眼相机" if args.fisheye else "普通相机"
        method_type = "映射方法" if use_remap else "直接方法"
        cv2.putText(comparison, f"类型: {camera_type} | 方法: {method_type}", 
                   (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
        cv2.putText(comparison, "按 's' 切换方法 | 按 'q' 退出", 
                   (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
        
        cv2.imshow(window_name, comparison)
        
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break
        elif key == ord('s'):
            use_remap = not use_remap
            method_name = "映射方法" if use_remap else "直接方法"
            print(f"切换到 {method_name}")
    
    cap.release()
    cv2.destroyAllWindows()
    print("测试完成")

if __name__ == "__main__":
    main()