# coding=utf-8

"""
眼在手上：使用采集的图片信息和机械臂位姿信息计算
相机坐标系相对于机械臂末端坐标系的旋转矩阵和平移向量
"""

import os
import logging
import yaml
import cv2
import numpy as np
from scipy.spatial.transform import Rotation as R

np.set_printoptions(precision=8, suppress=True)


current_path = os.path.dirname(os.path.abspath(__file__))
images_path = os.path.join(current_path, "data/20250804-1331")

# 加载标定板参数
with open("./config/config.yaml", 'r', encoding='utf-8') as file:
    data = yaml.safe_load(file)

XX = data.get("checkerboard_args").get("XX")
YY = data.get("checkerboard_args").get("YY")
L = data.get("checkerboard_args").get("L")

def gen_T_txt():
    joint_record_file = os.path.join(images_path, "joints_record.txt")
    with open(joint_record_file, "r") as f:
        lines = f.readlines()
    for line in lines:
        temp = line.replace("\n", "").split(" ")
        index = temp[0]
        joints = temp[2:]
        cmd = " ".join(["fk_example", "right"] + joints) + f" -o ./temp/T_ee_{index}.txt"
        print(cmd)
        os.system(cmd)
    print("gen_T_txt finished")

def load_robot_poses_from_fk(folder):
    """
    从fk_example输出的多个T_ee_{index}.txt文件构建RobotToolPose矩阵
    输出: 3x(4*N) numpy矩阵
    """
    files = [f for f in os.listdir(folder) if f.startswith("T_ee_") and f.endswith(".txt")]
    if not files:
        raise RuntimeError("未找到任何 T_ee_xxx.txt 文件")

    # 根据index排序
    files.sort(key=lambda x: int(x.split('_')[-1].split('.')[0]))

    poses = []
    for fname in files:
        path = os.path.join(folder, fname)
        data = np.loadtxt(path, delimiter=',')
        if data.shape != (12,):
            raise RuntimeError(f"{fname} 格式错误，期望12个值(r11,r12,r13,tx,...,tz)")

        # 3x4矩阵
        mat = data.reshape(3, 4)
        poses.append(mat)

    # 拼接成 3x(4*N)
    robot_pose_matrix = np.hstack(poses)
    out_path = os.path.join(folder, "../RobotToolPose.csv")
    np.savetxt(out_path, robot_pose_matrix, delimiter=',', fmt='%.8f')
    print(f"RobotToolPose.csv 已生成: {out_path}")
    return out_path


def func():
    path = os.path.dirname(__file__)
    # 设置寻找亚像素角点的参数
    criteria = (cv2.TERM_CRITERIA_MAX_ITER | cv2.TERM_CRITERIA_EPS, 30, 0.001)

    # 获取标定板角点的位置
    objp = np.zeros((XX * YY, 3), np.float32)
    objp[:, :2] = np.mgrid[0:XX, 0:YY].T.reshape(-1, 2)
    objp = L * objp

    obj_points = []     # 存储3D点
    img_points = []     # 存储2D点

    images_num = [f for f in os.listdir(images_path) if f.endswith('.png')]

    for i in range(1, len(images_num) + 1):
        image_file = os.path.join(images_path, f"img_{i}.png")
        if os.path.exists(image_file):
            print(f'读 {image_file}')
            img = cv2.imread(image_file)
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

            size = gray.shape[::-1]
            ret, corners = cv2.findChessboardCorners(gray, (XX, YY), None)

            if ret:
                cv2.drawChessboardCorners(img, (XX, YY), corners, ret)

                # 显示图像窗口
                save_path = os.path.join(current_path, "temp", f"res_{i}.jpg")
                cv2.imwrite(save_path, img)

                obj_points.append(objp)
                corners2 = cv2.cornerSubPix(gray, corners, (5, 5), (-1, -1), criteria)
                img_points.append(corners2 if corners2 is not None else corners)

    N = len(img_points)

    # 标定得到图案在相机坐标系下的位姿
    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(obj_points, img_points, size, None, None)

    print("-----------------------------------------------------")

    # === 新增：读取fk_example结果构建RobotToolPose ===
    csv_file = load_robot_poses_from_fk(os.path.join(path, "temp"))
    tool_pose = np.loadtxt(csv_file, delimiter=',')

    R_tool = []
    t_tool = []
    for i in range(int(N)):
        R_tool.append(tool_pose[0:3, 4*i:4*i+3])
        t_tool.append(tool_pose[0:3, 4*i+3])

    R_, t_ = cv2.calibrateHandEye(R_tool, t_tool, rvecs, tvecs, cv2.CALIB_HAND_EYE_TSAI)

    return R_, t_


if __name__ == '__main__':
    gen_T_txt()
    rotation_matrix, translation_vector = func()

    rotation = R.from_matrix(rotation_matrix)
    quaternion = rotation.as_quat()

    print(f"旋转矩阵是:\n {rotation_matrix}")
    print(f"平移向量是:\n {translation_vector}")
    print(f"四元数是：\n {quaternion}")