import cv2
import numpy as np
import glob
import os
from pathlib import Path

# ==================== 配置参数 ====================
CHESSBOARD_SIZE = (8, 11)  # 内角点数量 (列, 行)
SQUARE_SIZE = 20.0  # 每个方格实际边长（毫米），请改成你的真实尺寸

CAMERA_INDEX = 1
SAVE_DIR = "stereo_calib_images"  # 保存标定图片的文件夹
os.makedirs(SAVE_DIR, exist_ok=True)

# 标定用分辨率（每个单目）
CALIB_WIDTH = 1280
CALIB_HEIGHT = 720


# =================================================

def capture_images():
    """采集标定图片（左右目同步保存）"""
    cap = cv2.VideoCapture(CAMERA_INDEX, cv2.CAP_DSHOW)
    if not cap.isOpened():
        cap = cv2.VideoCapture(CAMERA_INDEX, cv2.CAP_MSMF)

    cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG'))
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 2560)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

    print("按空格键保存当前帧，按 q 退出采集")
    count = 0

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

        left = frame[:, :1280].copy()
        right = frame[:, 1280:].copy()

        gray_l = cv2.cvtColor(left, cv2.COLOR_BGR2GRAY)
        gray_r = cv2.cvtColor(right, cv2.COLOR_BGR2GRAY)

        ret_l, corners_l = cv2.findChessboardCorners(gray_l, CHESSBOARD_SIZE, None)
        ret_r, corners_r = cv2.findChessboardCorners(gray_r, CHESSBOARD_SIZE, None)

        # 用于显示的图像（带角点）
        display_left = left.copy()
        display_right = right.copy()

        if ret_l:
            cv2.drawChessboardCorners(display_left, CHESSBOARD_SIZE, corners_l, ret_l)
        if ret_r:
            cv2.drawChessboardCorners(display_right, CHESSBOARD_SIZE, corners_r, ret_r)

        display = np.hstack([display_left, display_right])
        cv2.putText(display, f"Saved: {count}  |  Space: Save  Q: Quit", (20, 40),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.imshow("Stereo Capture (Left | Right)", display)

        key = cv2.waitKey(1) & 0xFF
        if key == ord(' '):  # 空格保存
            if ret_l and ret_r:
                # 1. 保存干净的原图（用于正式标定）
                cv2.imwrite(f"{SAVE_DIR}/left_{count:02d}.png", left)
                cv2.imwrite(f"{SAVE_DIR}/right_{count:02d}.png", right)

                # 2. 同时保存带检测结果的可视化图（方便检查）
                cv2.imwrite(f"{SAVE_DIR}/left_{count:02d}_detect.png", display_left)
                cv2.imwrite(f"{SAVE_DIR}/right_{count:02d}_detect.png", display_right)

                print(f"已保存第 {count} 组（含检测结果图）")
                count += 1
            else:
                print("未同时检测到棋盘格，未保存")
        elif key == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()
    print(f"共采集 {count} 组图片")


def calibrate_stereo():
    """执行双目标定"""
    # 准备物体点（世界坐标）
    objp = np.zeros((CHESSBOARD_SIZE[0] * CHESSBOARD_SIZE[1], 3), np.float32)
    objp[:, :2] = np.mgrid[0:CHESSBOARD_SIZE[0], 0:CHESSBOARD_SIZE[1]].T.reshape(-1, 2)
    objp *= SQUARE_SIZE

    objpoints = []  # 3D点
    imgpoints_l = []  # 左目2D点
    imgpoints_r = []  # 右目2D点

    left_images = sorted(glob.glob(f"{SAVE_DIR}/left_*.png"))
    right_images = sorted(glob.glob(f"{SAVE_DIR}/right_*.png"))

    if len(left_images) < 10:
        print("标定图片太少（建议 ≥15 组），请重新采集")
        return

    print(f"开始处理 {len(left_images)} 组图片...")

    for left_path, right_path in zip(left_images, right_images):
        img_l = cv2.imread(left_path)
        img_r = cv2.imread(right_path)

        gray_l = cv2.cvtColor(img_l, cv2.COLOR_BGR2GRAY)
        gray_r = cv2.cvtColor(img_r, cv2.COLOR_BGR2GRAY)

        ret_l, corners_l = cv2.findChessboardCorners(gray_l, CHESSBOARD_SIZE, None)
        ret_r, corners_r = cv2.findChessboardCorners(gray_r, CHESSBOARD_SIZE, None)

        if ret_l and ret_r:
            # 亚像素精化
            criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
            corners_l = cv2.cornerSubPix(gray_l, corners_l, (11, 11), (-1, -1), criteria)
            corners_r = cv2.cornerSubPix(gray_r, corners_r, (11, 11), (-1, -1), criteria)

            objpoints.append(objp)
            imgpoints_l.append(corners_l)
            imgpoints_r.append(corners_r)
            print(f"✓ {Path(left_path).name}")
        else:
            print(f"✗ 跳过 {Path(left_path).name}")

    if len(objpoints) < 10:
        print("有效图片不足，标定失败")
        return

    # 单目标定（可选，用来给 stereoCalibrate 提供初始值）
    ret_l, mtx_l, dist_l, rvecs_l, tvecs_l = cv2.calibrateCamera(
        objpoints, imgpoints_l, (CALIB_WIDTH, CALIB_HEIGHT), None, None)
    ret_r, mtx_r, dist_r, rvecs_r, tvecs_r = cv2.calibrateCamera(
        objpoints, imgpoints_r, (CALIB_WIDTH, CALIB_HEIGHT), None, None)

    print("\n左目重投影误差:", ret_l)
    print("右目重投影误差:", ret_r)

    # 双目标定
    flags = cv2.CALIB_FIX_INTRINSIC  # 使用上面得到的内参
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 1e-5)

    ret, mtx_l, dist_l, mtx_r, dist_r, R, T, E, F = cv2.stereoCalibrate(
        objpoints, imgpoints_l, imgpoints_r,
        mtx_l, dist_l, mtx_r, dist_r,
        (CALIB_WIDTH, CALIB_HEIGHT),
        criteria=criteria,
        flags=flags
    )

    print("\n双目标定完成！")
    print("重投影误差:", ret)
    print("旋转矩阵 R:\n", R)
    print("平移向量 T:\n", T)

    # 计算立体校正参数
    R1, R2, P1, P2, Q, roi_l, roi_r = cv2.stereoRectify(
        mtx_l, dist_l, mtx_r, dist_r,
        (CALIB_WIDTH, CALIB_HEIGHT), R, T,
        alpha=0  # 0=只保留有效区域，1=保留所有像素
    )

    # 保存所有结果
    np.savez("stereo_calib_params.npz",
             mtx_l=mtx_l, dist_l=dist_l,
             mtx_r=mtx_r, dist_r=dist_r,
             R=R, T=T, E=E, F=F,
             R1=R1, R2=R2, P1=P1, P2=P2, Q=Q)

    print("\n标定参数已保存到 stereo_calib_params.npz")
    return mtx_l, dist_l, mtx_r, dist_r, R, T, Q


if __name__ == "__main__":
    print("1. 先采集标定图片")
    capture_images()

    print("\n2. 开始标定")
    calibrate_stereo()