using PyCall

# 引入 Python 的 OpenCV 和其他必要的库
cv2 = pyimport("cv2")
np = pyimport("numpy")
glob = pyimport("glob")

# === 棋盘参数 ===
pattern_size = (8, 6)  # 内部角点列数8，行数6
square_size = 1.0
image_path = "data/*.jpg"  # 修改为 data 文件夹下的 .jpg 文件

# === 停止条件 ===
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

# === 准备 3D 物体点 ===
cols, rows = pattern_size
objp = np.zeros((rows * cols, 3), dtype=np.float32)

# 3D 坐标点按棋盘的列和行来排列
objp[:, 1] .= repeat(collect(0:cols-1), outer=rows)  # X方向：列数
objp[:, 2] .= repeat(collect(0:rows-1), inner=cols)  # Y方向：行数

# 通过乘以 square_size 进行尺度调整
objp *= square_size

# === 存储 3D 点和 2D 点 ===
objpoints = []  # 存储 3D 物体点
imgpoints = []  # 存储 2D 图像点

# === 加载图像路径 ===
images = glob.glob(image_path)  # 直接使用 glob 函数获取文件列表

if isempty(images)
    println("没有找到图像")
else
    for fname in images
        println("处理图像: $fname")
        img = cv2.imread(fname)
        
        if img === nothing
            println("加载失败")
            continue
        end

        # 转为灰度图像
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        # 检查图像的尺寸
        height, width = size(gray)  # 获取图像尺寸
        println("图像尺寸：", height, " x ", width)

        # 进行高斯模糊
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)

        # 寻找棋盘角点
        found, corners = cv2.findChessboardCorners(blurred, (cols, rows))

        if found
            println("找到角点")
            push!(objpoints, objp)

            # 亚像素角点优化
            corners2 = cv2.cornerSubPix(gray, cv2.UMat(corners), (11, 11), (-1, -1), criteria)
            corners2 = np.array(corners2.get(), dtype=np.float32)
            corners2 = reshape(corners2, (size(corners2, 1), 2))
            push!(imgpoints, corners2)

            # 转为 numpy 数组
            objpoints_np = np.array(objpoints, dtype=np.float32)
            imgpoints_np = np.array(imgpoints, dtype=np.float32)

            # 相机标定
            result = cv2.calibrateCamera(objpoints_np, imgpoints_np, 
                (size(gray, 2), size(gray, 1)), nothing, nothing)

            ret, mtx, dist, rvecs, tvecs = result

            println("相机标定完成")
            println("内参矩阵 K = \n", mtx)
            println("畸变参数 dist = ", dist)

            # 对图像进行畸变矫正
            undistorted = cv2.undistort(img, mtx, dist)

            # 定义坐标轴的点
            axis = np.float32([ 
                [0, 0, 0],  # 原点
                [2, 0, 0],  # X轴
                [0, 2, 0],  # Y轴
                [0, 0, -2]   # Z轴
            ])

            # 获取当前旋转矩阵
            R_result = cv2.Rodrigues(rvecs[end])
            R = R_result[1]

            println("旋转向量 rvecs = ")
            println(rvecs[end])

            println("旋转矩阵 R = ")
            println(R)

            # 获取位移向量
            println("位移向量 t = ")
            println(tvecs[end])

            # 不做旋转，直接使用当前的R和T向量
            rvec_new = rvecs[end]

            # 投影坐标轴点
            imgpts = cv2.projectPoints(axis, rvec_new, tvecs[end], mtx, dist)[1]

            # 打印投影点
            println("投影点：")
            println(imgpts)

            imgpts = reshape(np.squeeze(imgpts), (4, 2))

            # 绘制坐标轴的函数
            function draw_axes(img, corners, imgpts)
                img_umat = cv2.UMat(img)

                # 获取第一个角点，作为原点
                corner = Tuple(Int.(round.(corners[1, :])))

                # 绘制坐标轴
                img_umat = cv2.line(img_umat, corner, Tuple(Int.(round.(imgpts[2, :]))), (255, 0, 0), 3)  # X轴
                img_umat = cv2.line(img_umat, corner, Tuple(Int.(round.(imgpts[3, :]))), (0, 255, 0), 3)  # Y轴
                img_umat = cv2.line(img_umat, corner, Tuple(Int.(round.(imgpts[4, :]))), (0, 0, 255), 3)  # Z轴

                return img_umat.get()
            end

            # 在矫正图像上绘制坐标轴
            img_with_axes = draw_axes(undistorted, corners2, imgpts)

            # 保存结果到 result_pose 文件夹
            output_path = joinpath("result_pose", basename(fname))
            cv2.imwrite(output_path, img_with_axes)
            println("成功保存图像到: $output_path")

            # 显示带角点和坐标轴的图像
            cv2.imshow("Undistorted Image with Axes and Corners", img_with_axes)
            cv2.waitKey(0)

        else
            println("未找到角点")
        end
    end
end

cv2.destroyAllWindows()