import os
import numpy as np # type: ignore
import cv2 as cv # type: ignore
import glob

# ==============================================
# 第一部分：相机标定
# ==============================================

# 终止条件
criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)

# 存储对象点和图像点的数组
objpoints = []  # 3D空间点
imgpoints = []  # 2D图像点

# 尝试的棋盘格尺寸（从大到小尝试）
possible_sizes = [(10, 8), (9, 7), (8, 6), (7, 5)]

# 文件夹设置
source_folder = "data"
result_folder = "result"
undistorted_folder = "undistorted"

# 创建输出目录
os.makedirs(result_folder, exist_ok=True)
os.makedirs(undistorted_folder, exist_ok=True)

def clear_folder(folder_path):
    """清空文件夹"""
    if os.path.exists(folder_path):
        for file in os.listdir(folder_path):
            file_path = os.path.join(folder_path, file)
            try:
                if os.path.isfile(file_path) or os.path.islink(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    clear_folder(file_path)
                    os.rmdir(file_path)
            except Exception as e:
                print(f'无法删除 {file_path}. 原因: {e}')

# 清空结果文件夹
clear_folder(result_folder)
clear_folder(undistorted_folder)

# 设置输出文件夹
pose_folder = "result_pose"
os.makedirs(pose_folder, exist_ok=True)
clear_folder(pose_folder)

# 获取所有标定图像
images = glob.glob(os.path.join(source_folder, '*.jpg'))

if not images:
    print("错误: 源文件夹中没有找到.jpg图像")
    exit()

# 第一阶段：检测所有图像的棋盘格角点
for fname in images:
    img = cv.imread(fname)
    if img is None:
        print(f"无法读取图像: {fname}")
        continue

    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    found = False
    
    for chessboard_size in possible_sizes:
        # 增强棋盘格检测
        ret, corners = cv.findChessboardCorners(
            gray, chessboard_size,
            flags=cv.CALIB_CB_ADAPTIVE_THRESH + 
                 cv.CALIB_CB_NORMALIZE_IMAGE +
                 cv.CALIB_CB_FAST_CHECK +
                 cv.CALIB_CB_FILTER_QUADS
        )
        
        if ret:
            print(f"找到棋盘格角点: {fname} 尺寸 {chessboard_size}")
            found = True

            # 创建对象点
            objp = np.zeros((chessboard_size[0] * chessboard_size[1], 3), np.float32)
            objp[:, :2] = np.mgrid[0:chessboard_size[1], 0:chessboard_size[0]].T.reshape(-1, 2)

            # 亚像素级角点检测
            corners2 = cv.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
            
            objpoints.append(objp)
            imgpoints.append(corners2)

            # 可视化检测结果
            cv.drawChessboardCorners(img, chessboard_size, corners2, ret)
            result_path = os.path.join(result_folder, f"corners_{os.path.basename(fname)}")
            cv.imwrite(result_path, img)
            break

    if not found:
        print(f"未检测到棋盘格: {fname}")

# 第二阶段：相机标定
if len(objpoints) < 10:
    print(f"错误: 只有 {len(objpoints)} 张图像成功检测 (至少需要10张)")
    exit()

# 初始化相机矩阵为单位矩阵
h, w = gray.shape[::-1]
initial_camera_matrix = np.eye(3, dtype=np.float32)
initial_camera_matrix[0, 0] = w  # 初始焦距 fx
initial_camera_matrix[1, 1] = h  # 初始焦距 fy
initial_camera_matrix[0, 2] = w / 2  # 主点 cx
initial_camera_matrix[1, 2] = h / 2  # 主点 cy

# 改进的标定参数
flags = (cv.CALIB_FIX_K3 + cv.CALIB_FIX_PRINCIPAL_POINT + 
         cv.CALIB_ZERO_TANGENT_DIST + cv.CALIB_USE_INTRINSIC_GUESS)

ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(
    objpoints, imgpoints, gray.shape[::-1], initial_camera_matrix, None,
    flags=flags
)

if not ret:
    print("相机标定失败!")
    exit()

print("\n标定成功!")
print("原始畸变系数:", dist)

# 畸变系数修正
if np.any(np.abs(dist) > 1.0):
    print("警告: 畸变系数异常，正在修正...")
    dist = np.clip(dist, -0.5, 0.5).astype(np.float32)
    print("修正后畸变系数:", dist)

# 第三阶段：图像去畸变
test_img = cv.imread(images[0])
if test_img is None:
    print("无法读取测试图像")
    exit()

h, w = test_img.shape[:2]
newcameramtx, roi = cv.getOptimalNewCameraMatrix(mtx, dist, (w, h), 0.8, (w, h))

def undistort_image(img):
    """改进的去畸变函数"""
    dst = cv.undistort(img, mtx, dist, None, newcameramtx)
    x, y, w, h = roi
    if w > 0 and h > 0:
        return dst[y:y+h, x:x+w]
    return dst

# 处理所有图像
for fname in images:
    img = cv.imread(fname)
    if img is None:
        continue
        
    dst = undistort_image(img)
    
    # 保存去畸变图像
    undistorted_path = os.path.join(undistorted_folder, os.path.basename(fname))
    cv.imwrite(undistorted_path, dst, [cv.IMWRITE_JPEG_QUALITY, 95])
    
    # 生成对比图
    comparison = np.hstack((cv.resize(img, (dst.shape[1], dst.shape[0])), dst))
    comparison_path = os.path.join(result_folder, f"compare_{os.path.basename(fname)}")
    cv.imwrite(comparison_path, comparison)

# 保存标定参数
np.savez(
    os.path.join(result_folder, "calibration_data.npz"),
    camera_matrix=mtx,
    distortion_coefficients=dist,
    optimal_matrix=newcameramtx
)

# 计算重投影误差
mean_error = 0
for i in range(len(objpoints)):
    imgpoints2, _ = cv.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)
    error = cv.norm(imgpoints[i], imgpoints2, cv.NORM_L2) / len(imgpoints2)
    mean_error += error

print(f"\n平均重投影误差: {mean_error/len(objpoints):.2f} 像素")
with open(os.path.join(result_folder, "calibration_report.txt"), "w") as f:
    f.write(f"平均重投影误差: {mean_error/len(objpoints):.2f} 像素\n")
    f.write(f"相机矩阵:\n{mtx}\n")
    f.write(f"畸变系数:\n{dist}\n")

print("标定完成! 结果保存在:", os.path.abspath(result_folder))

# ==============================================
# 第二部分：三维坐标系绘制
# ==============================================

# 加载标定数据
calib_data = np.load(os.path.join(result_folder, "calibration_data.npz"))
mtx = calib_data['camera_matrix']
dist = calib_data['distortion_coefficients']

def draw_axis(img, corners, imgpts):
    """绘制三维坐标系"""
    # 使用棋盘格的左下角作为坐标原点
    corner = tuple(corners[-1].ravel().astype(int))  # 左下角点
    imgpts = imgpts.astype(int)
    
    # 绘制X轴（蓝色）和Y轴（绿色）贴合棋盘边框
    img = cv.line(img, corner, tuple(imgpts[0].ravel()), (255, 0, 0), 5)  # X轴（蓝色）
    img = cv.line(img, corner, tuple(imgpts[1].ravel()), (0, 255, 0), 5)  # Y轴（绿色）
    img = cv.line(img, corner, tuple(imgpts[2].ravel()), (0, 0, 255), 5)  # Z轴（红色）
    return img

for fname in glob.glob(os.path.join(undistorted_folder, '*.jpg')):
    img = cv.imread(fname)
    if img is None:
        print(f"无法读取图像: {fname}")
        continue

    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    found = False
    
    # 尝试不同棋盘格尺寸
    for size in [(10, 8), (9, 7), (8, 6)]:
        ret, corners = cv.findChessboardCorners(gray, size, None)
        if ret:
            # 亚像素角点检测
            corners2 = cv.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
            
            # 创建对应的3D点
            objp = np.zeros((size[0] * size[1], 3), np.float32)
            objp[:, :2] = np.mgrid[0:size[1], 0:size[0]].T.reshape(-1, 2)
            
            # 计算位姿
            ret, rvecs, tvecs = cv.solvePnP(objp, corners2, mtx, dist)
            
            if ret:
                # 定义坐标系轴 (3个单位长度)
                axis = np.float32([[3, 0, 0], [0, 3, 0], [0, 0, -3]]).reshape(-1, 3)
                
                # 投影3D点到2D图像
                imgpts, _ = cv.projectPoints(axis, rvecs, tvecs, mtx, dist)
                
                # 绘制坐标系
                img = draw_axis(img, corners2, imgpts)
                
                # 保存结果
                output_path = os.path.join(pose_folder, os.path.basename(fname))
                if cv.imwrite(output_path, img):
                    print(f"成功保存坐标系图像: {output_path}")
                    found = True
                    break
    
    if not found:
        print(f"无法检测棋盘格或计算位姿: {fname}")

print("三维坐标系绘制完成! 结果保存在:", os.path.abspath(pose_folder))