#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
相机校准应用程序
使用检测到的棋盘格参数进行图像矫正和去畸变
"""

import cv2
import numpy as np
import json
import os


class CameraCalibration:
    def __init__(self, calibration_file='chessboard_detection_output/calibration_result.json'):
        """
        加载校准参数

        Args:
            calibration_file: 校准结果JSON文件路径
        """
        self.calibration_file = calibration_file
        self.camera_matrix = None
        self.dist_coeffs = None
        self.load_calibration()

    def load_calibration(self):
        """从JSON文件加载校准参数"""
        if not os.path.exists(self.calibration_file):
            raise FileNotFoundError(f"校准文件不存在: {self.calibration_file}")

        with open(self.calibration_file, 'r', encoding='utf-8') as f:
            data = json.load(f)

        self.camera_matrix = np.array(data['camera_matrix'])
        self.dist_coeffs = np.array(data['distortion_coefficients'])

        print("✓ 校准参数加载成功")
        print(f"  重投影误差: {data['reprojection_error']:.4f} 像素")
        print(f"  使用图像数: {data['num_images']}")

    def undistort_image(self, image):
        """
        对图像进行去畸变处理

        Args:
            image: 输入图像

        Returns:
            undistorted: 去畸变后的图像
        """
        h, w = image.shape[:2]
        new_camera_matrix, roi = cv2.getOptimalNewCameraMatrix(
            self.camera_matrix, self.dist_coeffs, (w, h), 1, (w, h)
        )

        # 去畸变
        undistorted = cv2.undistort(image, self.camera_matrix, self.dist_coeffs, None, new_camera_matrix)

        # 裁剪图像
        x, y, w, h = roi
        undistorted = undistorted[y:y+h, x:x+w]

        return undistorted

    def undistort_video(self, input_video, output_video='undistorted_video.avi'):
        """
        对视频进行去畸变处理

        Args:
            input_video: 输入视频路径
            output_video: 输出视频路径
        """
        cap = cv2.VideoCapture(input_video)
        if not cap.isOpened():
            print(f"无法打开视频: {input_video}")
            return

        # 获取视频参数
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

        # 读取第一帧获取尺寸
        ret, frame = cap.read()
        if not ret:
            print("无法读取视频帧")
            return

        h, w = frame.shape[:2]
        new_camera_matrix, roi = cv2.getOptimalNewCameraMatrix(
            self.camera_matrix, self.dist_coeffs, (w, h), 1, (w, h)
        )
        x, y, w_roi, h_roi = roi

        # 创建视频写入器
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        out = cv2.VideoWriter(output_video, fourcc, fps, (w_roi, h_roi))

        print(f"开始处理视频 (共 {total_frames} 帧)...")

        # 重置到开头
        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        frame_count = 0

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            frame_count += 1

            # 去畸变
            undistorted = cv2.undistort(frame, self.camera_matrix, self.dist_coeffs, None, new_camera_matrix)
            undistorted = undistorted[y:y+h_roi, x:x+w_roi]

            out.write(undistorted)

            if frame_count % 50 == 0:
                print(f"  处理进度: {frame_count}/{total_frames} ({100*frame_count/total_frames:.1f}%)")

        cap.release()
        out.release()

        print(f"\n✓ 视频处理完成，已保存到: {output_video}")

    def compare_images(self, image_path, output_path='comparison.jpg'):
        """
        生成原始图像和去畸变图像的对比图

        Args:
            image_path: 输入图像路径
            output_path: 输出对比图路径
        """
        image = cv2.imread(image_path)
        if image is None:
            print(f"无法读取图像: {image_path}")
            return

        undistorted = self.undistort_image(image)

        # 调整尺寸以便并排显示
        h1, w1 = image.shape[:2]
        h2, w2 = undistorted.shape[:2]
        h = min(h1, h2)
        image_resized = cv2.resize(image, (int(w1 * h / h1), h))
        undistorted_resized = cv2.resize(undistorted, (int(w2 * h / h2), h))

        # 并排拼接
        comparison = np.hstack([image_resized, undistorted_resized])

        # 添加文字标注
        cv2.putText(comparison, 'Original', (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 0, 255), 3)
        cv2.putText(comparison, 'Undistorted', (w1 + 50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 0), 3)

        cv2.imwrite(output_path, comparison)
        print(f"✓ 对比图已保存到: {output_path}")


def main():
    """主函数 - 演示如何使用校准参数"""
    print("=== 相机校准应用程序 ===\n")

    # 加载校准参数
    calib = CameraCalibration()

    # 示例1: 对检测结果图像进行去畸变
    output_dir = 'chessboard_detection_output'
    detected_images = [f for f in os.listdir(output_dir) if f.startswith('detected_') and f.endswith('.jpg')]

    if detected_images:
        print(f"\n找到 {len(detected_images)} 张检测图像")
        sample_image = os.path.join(output_dir, detected_images[0])
        print(f"生成对比图: {sample_image}")
        calib.compare_images(sample_image, os.path.join(output_dir, 'comparison.jpg'))

    # 示例2: 对原始视频进行去畸变
    print("\n是否要对原始视频进行去畸变处理？")
    print("注意: 这将处理整个视频，可能需要一些时间")
    print("如需处理，请取消注释下面的代码行：")
    print("# calib.undistort_video('Video_20260303114232727.avi', 'undistorted_video.avi')")


if __name__ == '__main__':
    main()