#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
棋盘格检测程序
用于检测视频中的11x8内角点棋盘格，并用于相机校准
"""

import cv2
import numpy as np
import os
import json
from datetime import datetime


class ChessboardDetector:
    def __init__(self, pattern_size=(11, 8), square_size=1.0):
        """
        初始化棋盘格检测器

        Args:
            pattern_size: 棋盘格内角点数量 (列, 行)
            square_size: 棋盘格方格实际尺寸(单位:mm或其他)
        """
        self.pattern_size = pattern_size
        self.square_size = square_size

        # 准备棋盘格的3D坐标点
        self.objp = np.zeros((pattern_size[0] * pattern_size[1], 3), np.float32)
        self.objp[:, :2] = np.mgrid[0:pattern_size[0], 0:pattern_size[1]].T.reshape(-1, 2)
        self.objp *= square_size

        # 存储所有图像的角点
        self.obj_points = []  # 3D点
        self.img_points = []  # 2D点

        # 角点检测参数
        self.criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

    def detect_chessboard(self, image):
        """
        检测单张图像中的棋盘格

        Args:
            image: 输入图像

        Returns:
            ret: 是否检测成功
            corners: 角点坐标
            gray: 灰度图
        """
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # 查找棋盘格角点
        ret, corners = cv2.findChessboardCorners(gray, self.pattern_size, None)

        if ret:
            # 亚像素精度优化
            corners = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), self.criteria)

        return ret, corners, gray

    def process_video(self, video_path, output_dir='output', sample_interval=30):
        """
        处理视频文件，检测棋盘格

        Args:
            video_path: 视频文件路径
            output_dir: 输出目录
            sample_interval: 采样间隔(帧数)
        """
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            print(f"无法打开视频: {video_path}")
            return False

        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = cap.get(cv2.CAP_PROP_FPS)

        print(f"视频信息: 总帧数={total_frames}, FPS={fps}")
        print(f"开始检测棋盘格 (内角点: {self.pattern_size[0]}x{self.pattern_size[1]})...")

        frame_count = 0
        detected_count = 0

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

            frame_count += 1

            # 按间隔采样
            if frame_count % sample_interval != 0:
                continue

            # 检测棋盘格
            success, corners, gray = self.detect_chessboard(frame)

            if success:
                detected_count += 1
                self.obj_points.append(self.objp)
                self.img_points.append(corners)

                # 绘制角点
                vis_img = frame.copy()
                cv2.drawChessboardCorners(vis_img, self.pattern_size, corners, success)

                # 保存结果图像
                output_path = os.path.join(output_dir, f'detected_{detected_count:03d}_frame{frame_count}.jpg')
                cv2.imwrite(output_path, vis_img)

                print(f"✓ 帧 {frame_count}: 检测成功 (已保存 {detected_count} 张)")
            else:
                print(f"✗ 帧 {frame_count}: 未检测到棋盘格")

        cap.release()

        print(f"\n检测完成: 共处理 {frame_count} 帧, 成功检测 {detected_count} 张图像")
        return detected_count > 0

    def calibrate_camera(self, image_size):
        """
        执行相机校准

        Args:
            image_size: 图像尺寸 (width, height)

        Returns:
            ret: 标定误差
            camera_matrix: 相机内参矩阵
            dist_coeffs: 畸变系数
            rvecs: 旋转向量
            tvecs: 平移向量
        """
        if len(self.obj_points) < 3:
            print("错误: 需要至少3张成功检测的图像进行校准")
            return None

        print(f"\n开始相机校准 (使用 {len(self.obj_points)} 张图像)...")

        ret, camera_matrix, dist_coeffs, rvecs, tvecs = cv2.calibrateCamera(
            self.obj_points, self.img_points, image_size, None, None
        )

        print(f"标定完成! 重投影误差: {ret:.4f} 像素")

        return ret, camera_matrix, dist_coeffs, rvecs, tvecs

    def save_calibration_results(self, camera_matrix, dist_coeffs, ret, output_path='calibration_result.json'):
        """
        保存校准结果到JSON文件
        """
        result = {
            'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            'pattern_size': self.pattern_size,
            'square_size': self.square_size,
            'num_images': len(self.obj_points),
            'reprojection_error': float(ret),
            'camera_matrix': camera_matrix.tolist(),
            'distortion_coefficients': dist_coeffs.tolist()
        }

        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(result, f, indent=4, ensure_ascii=False)

        print(f"\n校准结果已保存到: {output_path}")

        # 打印结果
        print("\n=== 相机校准结果 ===")
        print(f"重投影误差: {ret:.4f} 像素")
        print(f"\n相机内参矩阵:")
        print(camera_matrix)
        print(f"\n畸变系数:")
        print(dist_coeffs.ravel())


def main():
    # 配置参数
    VIDEO_PATH = 'Video_20260303114232727.avi'
    OUTPUT_DIR = 'chessboard_detection_output'
    PATTERN_SIZE = (11, 8)  # 11x8 内角点
    SQUARE_SIZE = 25.0  # 假设每个方格25mm，根据实际情况调整
    SAMPLE_INTERVAL = 30  # 每30帧采样一次

    # 创建检测器
    detector = ChessboardDetector(pattern_size=PATTERN_SIZE, square_size=SQUARE_SIZE)

    # 处理视频
    success = detector.process_video(VIDEO_PATH, OUTPUT_DIR, SAMPLE_INTERVAL)

    if not success:
        print("未能检测到任何棋盘格，程序退出")
        return

    # 获取图像尺寸
    cap = cv2.VideoCapture(VIDEO_PATH)
    ret, frame = cap.read()
    if ret:
        image_size = (frame.shape[1], frame.shape[0])
        cap.release()

        # 执行相机校准
        calib_result = detector.calibrate_camera(image_size)

        if calib_result:
            ret, camera_matrix, dist_coeffs, rvecs, tvecs = calib_result

            # 保存校准结果
            detector.save_calibration_results(
                camera_matrix, dist_coeffs, ret,
                os.path.join(OUTPUT_DIR, 'calibration_result.json')
            )

            print("\n✓ 校准完成！可以使用生成的校准参数进行图像矫正")


if __name__ == '__main__':
    main()