#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
可视化视频中的棋盘格检测
实时显示每一帧的检测结果
"""

import cv2
import numpy as np


class ChessboardVisualizer:
    def __init__(self, pattern_size=(11, 8)):
        """
        初始化可视化器

        Args:
            pattern_size: 棋盘格内角点数量 (列, 行)
        """
        self.pattern_size = pattern_size
        self.criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

    def visualize_video(self, video_path, output_path='visualized_video.avi', show_window=False):
        """
        可视化整个视频的棋盘格检测

        Args:
            video_path: 输入视频路径
            output_path: 输出视频路径
            show_window: 是否显示实时窗口
        """
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            print(f"无法打开视频: {video_path}")
            return

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

        print(f"视频信息: {width}x{height}, {fps:.2f} FPS, {total_frames} 帧")

        # 创建视频写入器
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

        frame_count = 0
        detected_count = 0

        print(f"\n开始处理视频...")
        print("绿色角点 = 检测成功, 红色文字 = 未检测到")

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

            frame_count += 1

            # 转换为灰度图
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

            # 检测棋盘格
            ret_detect, corners = cv2.findChessboardCorners(gray, self.pattern_size, None)

            # 创建可视化图像
            vis_frame = frame.copy()

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

                # 绘制角点
                cv2.drawChessboardCorners(vis_frame, self.pattern_size, corners, ret_detect)

                # 添加成功标记
                cv2.putText(vis_frame, f'DETECTED #{detected_count}', (20, 40),
                           cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            else:
                # 添加未检测标记
                cv2.putText(vis_frame, 'NOT DETECTED', (20, 40),
                           cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

            # 添加帧信息
            cv2.putText(vis_frame, f'Frame: {frame_count}/{total_frames}', (20, height - 20),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

            # 写入输出视频
            out.write(vis_frame)

            # 显示窗口（可选）
            if show_window:
                cv2.imshow('Chessboard Detection', vis_frame)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    print("\n用户中断")
                    break

            # 进度显示
            if frame_count % 50 == 0:
                progress = 100 * frame_count / total_frames
                print(f"  进度: {frame_count}/{total_frames} ({progress:.1f}%) - 已检测: {detected_count}")

        cap.release()
        out.release()
        if show_window:
            cv2.destroyAllWindows()

        print(f"\n✓ 处理完成!")
        print(f"  总帧数: {frame_count}")
        print(f"  检测成功: {detected_count} 帧 ({100*detected_count/frame_count:.1f}%)")
        print(f"  输出文件: {output_path}")


def main():
    VIDEO_PATH = 'Video_20260303114232727.avi'
    OUTPUT_PATH = 'visualized_chessboard_detection.avi'
    PATTERN_SIZE = (11, 8)

    visualizer = ChessboardVisualizer(pattern_size=PATTERN_SIZE)
    visualizer.visualize_video(VIDEO_PATH, OUTPUT_PATH, show_window=False)


if __name__ == '__main__':
    main()