import numpy as np
import cv2
from typing import List, Tuple
import argparse
import json
import os
import glob
import csv

def hand_eye_calibration(robot_poses: List[np.ndarray],
                         camera_poses: List[np.ndarray],
                         mode: str = 'eye_in_hand') -> np.ndarray:
    """
    执行手眼标定，解决AX = XB问题。

    参数:
        robot_poses (list of np.array): 机械臂末端在基座坐标系中的位姿列表 (4x4齐次矩阵)。
        camera_poses (list of np.array): 标定板在相机坐标系中的位姿列表 (4x4齐次矩阵)。注意：这是 T_c_t（target->camera），与OpenCV的 R_target2cam/t_target2cam 一致。
        mode (str): 标定模式。'eye_in_hand' 或 'eye_to_hand'。

    返回:
        X (np.array): 求解出的变换矩阵X (4x4齐次矩阵)。
                      对于 eye_in_hand: X = camera^T_end_effector (相机在机械臂末端坐标系中的位姿)
                      对于 eye_to_hand: X = base^T_camera (相机在机器人基座坐标系中的位姿)
    """
    # 输入验证
    n = len(robot_poses)
    if n != len(camera_poses):
        raise ValueError("机器人位姿数量与相机位姿数量必须相同。")
    if n < 3:
        raise ValueError("至少需要三组数据（且包含显著旋转）才能进行稳定标定。")

    # OpenCV calibrateHandEye 需要每一帧的绝对位姿：
    #  - R_gripper2base, t_gripper2base: T_g^b（夹爪到基座）
    #  - R_target2cam,  t_target2cam:  T_t^c（标定板到相机）
    # 传入时请确保 robot_poses 为 T_b^g（基座到夹爪），camera_poses 为 T_c^t（目标到相机）
    R_gripper2base: List[np.ndarray] = []
    t_gripper2base: List[np.ndarray] = []
    R_target2cam: List[np.ndarray] = []
    t_target2cam: List[np.ndarray] = []

    for i in range(n):
        T_b_g = robot_poses[i]
        if T_b_g.shape != (4, 4):
            raise ValueError("robot_poses 中的矩阵必须是 4x4")
        T_g_b = np.linalg.inv(T_b_g)
        R_gripper2base.append(T_g_b[:3, :3])
        t_gripper2base.append(T_g_b[:3, 3].reshape(3, 1))

        T_c_t = camera_poses[i]
        if T_c_t.shape != (4, 4):
            raise ValueError("camera_poses 中的矩阵必须是 4x4")
        R_target2cam.append(T_c_t[:3, :3])
        t_target2cam.append(T_c_t[:3, 3].reshape(3, 1))

    # 选择算法，失败时自动尝试其他方法
    methods = [
        getattr(cv2, 'CALIB_HAND_EYE_TSAI', None),
        getattr(cv2, 'CALIB_HAND_EYE_PARK', None),
        getattr(cv2, 'CALIB_HAND_EYE_HORAUD', None),
        getattr(cv2, 'CALIB_HAND_EYE_ANDREFF', None),
        getattr(cv2, 'CALIB_HAND_EYE_DANIILIDIS', None),
    ]
    methods = [m for m in methods if m is not None]

    last_err = None
    for method in methods:
        try:
            R_cam2gripper, t_cam2gripper = cv2.calibrateHandEye(
                R_gripper2base, t_gripper2base,
                R_target2cam, t_target2cam,
                method=method,
            )
            # 成功返回
            X_cam_gripper = np.eye(4)
            X_cam_gripper[:3, :3] = R_cam2gripper
            X_cam_gripper[:3, 3] = t_cam2gripper.flatten()

            if mode == 'eye_in_hand':
                # 直接返回 camera^T_gripper
                return X_cam_gripper
            else:
                # eye_to_hand：需要 base^T_camera。已知：
                #  对每一帧 i，有 T_b^c_i = T_b^g_i * T_g^c（其中 T_g^c = (T_c^g)^{-1}）
                T_g_c = np.linalg.inv(X_cam_gripper)
                T_b_c_list = []
                for T_b_g in robot_poses:
                    T_b_c_list.append(T_b_g @ T_g_c)
                return average_SE3(T_b_c_list)
        except Exception as e:
            last_err = e
            continue

    raise RuntimeError(f"hand-eye 标定失败：{last_err}")

def create_homogeneous_matrix(rvec, tvec):
    """
    根据旋转向量和平移向量创建4x4齐次变换矩阵。

    参数:
        rvec: 3x1旋转向量
        tvec: 3x1平移向量

    返回:
        T: 4x4齐次变换矩阵
    """
    R, _ = cv2.Rodrigues(rvec)
    T = np.eye(4)
    T[:3, :3] = R
    T[:3, 3] = tvec.flatten()
    return T


def rot_to_quat(R: np.ndarray) -> np.ndarray:
    """将旋转矩阵转为四元数 [x, y, z, w]（右手，单位四元数）。"""
    tr = np.trace(R)
    if tr > 0.0:
        S = np.sqrt(tr + 1.0) * 2.0
        qw = 0.25 * S
        qx = (R[2, 1] - R[1, 2]) / S
        qy = (R[0, 2] - R[2, 0]) / S
        qz = (R[1, 0] - R[0, 1]) / S
    else:
        if R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
            S = np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2.0
            qx = 0.25 * S
            qy = (R[0, 1] + R[1, 0]) / S
            qz = (R[0, 2] + R[2, 0]) / S
            qw = (R[2, 1] - R[1, 2]) / S
        elif R[1, 1] > R[2, 2]:
            S = np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2.0
            qx = (R[0, 1] + R[1, 0]) / S
            qy = 0.25 * S
            qz = (R[1, 2] + R[2, 1]) / S
            qw = (R[0, 2] - R[2, 0]) / S
        else:
            S = np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2.0
            qx = (R[0, 2] + R[2, 0]) / S
            qy = (R[1, 2] + R[2, 1]) / S
            qz = 0.25 * S
            qw = (R[1, 0] - R[0, 1]) / S
    q = np.array([qx, qy, qz, qw])
    return q / (np.linalg.norm(q) + 1e-12)


def quat_to_rot(q: np.ndarray) -> np.ndarray:
    """四元数 [x, y, z, w] 转旋转矩阵。"""
    x, y, z, w = q
    xx, yy, zz = x*x, y*y, z*z
    xy, xz, yz = x*y, x*z, y*z
    wx, wy, wz = w*x, w*y, w*z
    R = np.array([
        [1 - 2*(yy + zz), 2*(xy - wz),     2*(xz + wy)],
        [2*(xy + wz),     1 - 2*(xx + zz), 2*(yz - wx)],
        [2*(xz - wy),     2*(yz + wx),     1 - 2*(xx + yy)]
    ])
    return R


def average_SE3(T_list: List[np.ndarray]) -> np.ndarray:
    """对一组 SE(3) 变换做简单平均（四元数+平移平均）。"""
    if not T_list:
        raise ValueError("T_list 为空")
    # 平移均值
    t_stack = np.stack([T[:3, 3] for T in T_list], axis=0)
    t_mean = np.mean(t_stack, axis=0)
    # 旋转均值（单位四元数平均后归一化）
    q_list = [rot_to_quat(T[:3, :3]) for T in T_list]
    # 对齐四元数符号到第一个四元数的半球，避免相互抵消
    q0 = q_list[0]
    aligned = []
    for q in q_list:
        if np.dot(q, q0) < 0:
            aligned.append(-q)
        else:
            aligned.append(q)
    q = np.mean(np.stack(aligned, axis=0), axis=0)
    q = q / (np.linalg.norm(q) + 1e-12)
    R_mean = quat_to_rot(q)
    T = np.eye(4)
    T[:3, :3] = R_mean
    T[:3, 3] = t_mean
    return T


def hand_eye_residual(robot_poses: List[np.ndarray], camera_poses: List[np.ndarray], X_cam_gripper: np.ndarray) -> Tuple[float, float]:
    """基于 AX ≈ X B 的关系计算残差（旋转角度均值，平移均值）。
    返回 (rot_err_deg, trans_err_m)。"""
    rot_errs = []
    trans_errs = []
    n = len(robot_poses)
    # 构造相对运动对，避免偏置
    for i in range(n - 1):
        A = np.linalg.inv(robot_poses[i + 1]) @ robot_poses[i]  # T_b^g_{i+1}^{-1} * T_b^g_i
        B = np.linalg.inv(camera_poses[i + 1]) @ camera_poses[i]  # T_c^t_{i+1}^{-1} * T_c^t_i
        # OpenCV定义 X 为 T_c^g（camera->gripper）
        lhs = np.linalg.inv(A)  # OpenCV论文中常用 A 为 g2b，但这里用的相对位姿定义略有差异
        # 直接用 AX 与 XB 的等价残差度量
        AX = A @ np.linalg.inv(X_cam_gripper)
        XB = np.linalg.inv(X_cam_gripper) @ B
        Delta = np.linalg.inv(AX) @ XB
        R = Delta[:3, :3]
        t = Delta[:3, 3]
        # 旋转角度
        angle = np.rad2deg(np.arccos(np.clip((np.trace(R) - 1) / 2.0, -1.0, 1.0)))
        rot_errs.append(abs(angle))
        trans_errs.append(np.linalg.norm(t))
    return float(np.mean(rot_errs)), float(np.mean(trans_errs))

# ==================== 示例用法 ====================

def generate_synthetic_dataset(num_poses: int = 12, seed: int = 42) -> Tuple[List[np.ndarray], List[np.ndarray]]:
    """生成可解的模拟数据集 (T_b^g 列表, T_c^t 列表)。"""
    rng = np.random.default_rng(seed)

    def hat(v):
        return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])

    def exp_so3(w):
        th = np.linalg.norm(w)
        if th < 1e-12:
            return np.eye(3)
        k = w / th
        K = hat(k)
        return np.eye(3) + np.sin(th) * K + (1 - np.cos(th)) * (K @ K)

    def make_T(R, t):
        T = np.eye(4)
        T[:3, :3] = R
        T[:3, 3] = t
        return T

    # 真实（未知）相机在夹爪中的位姿：T_c^g
    R_c_g = exp_so3(np.deg2rad([10, -5, 15]))
    t_c_g = np.array([0.03, -0.02, 0.15])
    T_c_g_true = make_T(R_c_g, t_c_g)

    # 世界（=基座）到标定板的固定位姿：T_b^t
    R_b_t = exp_so3(np.deg2rad([0, 0, 0]))
    t_b_t = np.array([0.4, 0.0, 0.3])
    T_b_t = make_T(R_b_t, t_b_t)

    robot_poses: List[np.ndarray] = []  # T_b^g
    camera_poses: List[np.ndarray] = [] # T_c^t

    for _ in range(num_poses):
        # 生成多轴、较大幅度的运动（避免退化）
        ang = np.deg2rad(rng.uniform([-30, -30, -30], [30, 30, 30]))
        # 让 SO(3) 采样更丰富：
        R_b_g = (np.eye(3) @ rot_perturb(ang))
        t_b_g = rng.uniform([-0.2, -0.2, 0.3], [0.2, 0.2, 0.6])
        T_b_g = make_T(R_b_g, t_b_g)
        robot_poses.append(T_b_g)

        # 相机位姿：T_b^c = T_b^g * T_g^c，其中 T_g^c = (T_c^g)^{-1}
        T_g_c = np.linalg.inv(T_c_g_true)
        T_b_c = T_b_g @ T_g_c

        # 目标在相机坐标系：T_c^t = (T_b^c)^{-1} * T_b^t
        T_c_t = np.linalg.inv(T_b_c) @ T_b_t
        camera_poses.append(T_c_t)

    return robot_poses, camera_poses


def rot_perturb(ang_vec: np.ndarray) -> np.ndarray:
    """给定欧拉角近似（rad）的向量，输出旋转矩阵（用罗德里格公式按轴角）。"""
    axis = ang_vec
    th = np.linalg.norm(axis)
    if th < 1e-12:
        return np.eye(3)
    k = axis / th
    K = np.array([[0, -k[2], k[1]], [k[2], 0, -k[0]], [-k[1], k[0], 0]])
    return np.eye(3) + np.sin(th) * K + (1 - np.cos(th)) * (K @ K)


def to_json_SE3(T: np.ndarray) -> dict:
    R = T[:3, :3]
    t = T[:3, 3]
    q = rot_to_quat(R)
    return {
        "matrix": T.tolist(),
        "rotation_matrix": R.tolist(),
        "translation": t.tolist(),
        "quaternion_xyzw": q.tolist(),
    }


def load_se3_matrix(path: str) -> np.ndarray:
    """加载 4x4 齐次矩阵。支持 .npy / .npz(json-like)/.json。"""
    if not os.path.exists(path):
        raise FileNotFoundError(path)
    ext = os.path.splitext(path)[1].lower()
    if ext == ".npy":
        M = np.load(path)
    elif ext == ".npz":
        d = np.load(path)
        key = 'matrix' if 'matrix' in d.files else d.files[0]
        M = d[key]
    elif ext == ".json":
        with open(path, 'r') as f:
            data = json.load(f)
        M = np.array(data.get('matrix', data))
    else:
        raise ValueError(f"不支持的文件类型: {ext}")
    M = np.array(M)
    if M.shape != (4, 4):
        raise ValueError(f"矩阵形状应为(4,4)，实际为{M.shape}")
    return M


def load_camera_calib(path: str) -> Tuple[np.ndarray, np.ndarray]:
    """加载相机内参，支持 OpenCV YAML/JSON（cameraMatrix, distCoeffs）或简单 JSON：{"K":[], "D":[]}。"""
    import yaml
    with open(path, 'r') as f:
        txt = f.read()
    data = None
    try:
        data = yaml.safe_load(txt)
    except Exception:
        try:
            data = json.loads(txt)
        except Exception:
            raise ValueError("无法解析相机标定文件，支持 YAML/JSON，包含 camera_matrix/cameraMatrix/K 和 dist_coeffs/distCoeffs/D 字段")

    # 支持多种键名
    K = data.get('camera_matrix') or data.get('cameraMatrix') or data.get('K')
    if isinstance(K, dict) and 'data' in K:
        K = K['data']
    D = data.get('dist_coeffs') or data.get('distCoeffs') or data.get('D')
    if isinstance(D, dict) and 'data' in D:
        D = D['data']
    if K is None or D is None:
        raise ValueError("标定文件需包含 camera_matrix/cameraMatrix/K 与 dist_coeffs/distCoeffs/D")
    K = np.array(K, dtype=float).reshape(3, 3)
    D = np.array(D, dtype=float).reshape(-1, 1)
    return K, D


def build_chessboard_object_points(rows: int, cols: int, square: float) -> np.ndarray:
    """生成棋盘格 3D 角点（单位：米），Z=0 平面，原点在棋盘左上角角点。"""
    objp = np.zeros((rows * cols, 3), np.float32)
    grid = np.mgrid[0:cols, 0:rows].T.reshape(-1, 2)
    objp[:, :2] = grid * square
    return objp


def detect_chessboard(gray: np.ndarray, rows: int, cols: int) -> Tuple[bool, np.ndarray]:
    flags = cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_NORMALIZE_IMAGE
    ok, corners = cv2.findChessboardCorners(gray, (cols, rows), flags)
    if not ok:
        return False, None
    # 亚像素优化
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    corners = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
    return True, corners.reshape(-1, 2)


def detect_charuco(gray: np.ndarray, rows: int, cols: int, square: float, marker: float, dict_name: str = 'DICT_4X4_50') -> Tuple[bool, np.ndarray, np.ndarray]:
    """检测 Charuco 角点，返回 (ok, charuco_corners, charuco_ids)。需要 opencv-contrib-python。"""
    if not hasattr(cv2, 'aruco'):
        raise RuntimeError("需要 opencv-contrib-python 才能使用 Charuco。请安装 opencv-contrib-python==4.11.0.86")
    aruco = cv2.aruco
    dictionary = getattr(aruco, dict_name)
    aruco_dict = aruco.getPredefinedDictionary(getattr(aruco, dict_name))
    board = aruco.CharucoBoard((cols, rows), square, marker, aruco_dict)
    params = aruco.DetectorParameters()
    detector = aruco.ArucoDetector(aruco_dict, params)
    corners, ids, _ = detector.detectMarkers(gray)
    if ids is None or len(ids) == 0:
        return False, None, None
    retval, charuco_corners, charuco_ids = aruco.interpolateCornersCharuco(corners, ids, gray, board)
    ok = retval is not None and retval >= 4
    return ok, (charuco_corners.reshape(-1, 2) if ok else None), (charuco_ids if ok else None)


def solve_pnp_from_image(image_path: str, K: np.ndarray, D: np.ndarray,
                         pattern: str, rows: int, cols: int, square: float,
                         charuco_marker: float = None) -> np.ndarray:
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise FileNotFoundError(f"无法读取图像: {image_path}")
    if pattern == 'chessboard':
        ok, corners = detect_chessboard(img, rows, cols)
        if not ok:
            raise RuntimeError(f"棋盘检测失败: {image_path}")
        objp = build_chessboard_object_points(rows, cols, square)
        # PnP
        retval, rvec, tvec = cv2.solvePnP(objp, corners, K, D, flags=cv2.SOLVEPNP_ITERATIVE)
        if not retval:
            raise RuntimeError(f"solvePnP 失败: {image_path}")
        return create_homogeneous_matrix(rvec, tvec)
    elif pattern == 'charuco':
        if charuco_marker is None:
            raise ValueError("使用 charuco 时必须提供 --charuco-marker（米）")
        ok, charuco_corners, charuco_ids = detect_charuco(img, rows, cols, square, charuco_marker)
        if not ok:
            raise RuntimeError(f"Charuco 检测失败: {image_path}")
        # 从 board 索引构建 3D 角点
        aruco = cv2.aruco
        aruco_dict = aruco.getPredefinedDictionary(getattr(aruco, 'DICT_4X4_50'))
        board = aruco.CharucoBoard((cols, rows), square, charuco_marker, aruco_dict)
        obj_pts = []
        img_pts = []
        for idx, corner in zip(charuco_ids.flatten(), charuco_corners.reshape(-1, 2)):
            obj_pts.append(board.getChessboardCorners()[idx][0])
            img_pts.append(corner)
        obj_pts = np.array(obj_pts, dtype=np.float32)
        img_pts = np.array(img_pts, dtype=np.float32)
        retval, rvec, tvec = cv2.solvePnP(obj_pts, img_pts, K, D, flags=cv2.SOLVEPNP_ITERATIVE)
        if not retval:
            raise RuntimeError(f"solvePnP 失败: {image_path}")
        return create_homogeneous_matrix(rvec, tvec)
    else:
        raise ValueError("pattern 仅支持 chessboard 或 charuco")


def load_robot_poses_file(path: str) -> List[np.ndarray]:
    """从文件加载 robot_poses 列表。支持 .npz(robot_poses), .json(list or dict.matrix), .csv( image, tx,ty,tz,qx,qy,qz,qw )。"""
    ext = os.path.splitext(path)[1].lower()
    if ext == '.npz':
        d = np.load(path)
        arr = d['robot_poses']
        assert arr.ndim == 3 and arr.shape[1:] == (4, 4)
        return [arr[i] for i in range(arr.shape[0])]
    if ext == '.json':
        with open(path, 'r') as f:
            data = json.load(f)
        if isinstance(data, dict) and 'robot_poses' in data:
            mats = data['robot_poses']
        else:
            mats = data
        return [np.array(M).reshape(4, 4) for M in mats]
    if ext == '.csv':
        poses = []
        with open(path, 'r') as f:
            reader = csv.DictReader(f)
            for row in reader:
                tx, ty, tz = float(row['tx']), float(row['ty']), float(row['tz'])
                qx, qy, qz, qw = float(row['qx']), float(row['qy']), float(row['qz']), float(row['qw'])
                R = quat_to_rot(np.array([qx, qy, qz, qw]))
                T = np.eye(4)
                T[:3, :3] = R
                T[:3, 3] = np.array([tx, ty, tz])
                poses.append(T)
        return poses
    raise ValueError(f"不支持的 robot poses 文件类型: {ext}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Hand-Eye calibration (eye-in-hand / eye-to-hand)")
    parser.add_argument("--data", type=str, default=None,
                        help=".npz dataset with arrays 'robot_poses' and 'camera_poses' of shape (N,4,4)")
    parser.add_argument("--mode", type=str, default=None, choices=["eye_in_hand", "eye_to_hand"],
                        help="Calibration mode. If omitted, both modes are computed")
    parser.add_argument("--out", type=str, default=None, help="Output JSON file for the calibration result")
    parser.add_argument("--export-dataset", type=str, default=None,
                        help="Export the dataset used (robot_poses,camera_poses) to .npz at this path")
    parser.add_argument("--seed", type=int, default=42, help="Synthetic dataset RNG seed when no --data is provided")
    parser.add_argument("--num-poses", type=int, default=12, help="Synthetic dataset pose count when no --data is provided")
    parser.add_argument("--eval", action="store_true", help="Compute simple AX≈XB residual metrics (eye-in-hand)")
    # Image-based camera pose estimation
    parser.add_argument("--images", type=str, default=None,
                        help="Directory or glob for images captured by the depth camera (color/IR). Sorted alphanumerically.")
    parser.add_argument("--camera-calib", type=str, default=None, help="Camera intrinsic file (YAML/JSON)")
    parser.add_argument("--pattern", type=str, default="chessboard", choices=["chessboard", "charuco"], help="Calibration board pattern")
    parser.add_argument("--board-rows", type=int, default=6, help="Inner rows of chessboard/charuco")
    parser.add_argument("--board-cols", type=int, default=9, help="Inner cols of chessboard/charuco")
    parser.add_argument("--square-size", type=float, default=0.024, help="Square size in meters")
    parser.add_argument("--charuco-marker", type=float, default=None, help="Charuco marker size in meters (for charuco only)")
    parser.add_argument("--robot-poses-file", type=str, default=None, help="Path to robot poses file (.npz/.json/.csv)")
    # Frames handling
    parser.add_argument("--robot-poses-frame", type=str, default="robot_base", choices=["robot_base", "arm_base"],
                        help="Frame of robot_poses in the dataset: robot_base (default) or arm_base")
    parser.add_argument("--arm-to-robot-base", type=str, default=None,
                        help="Optional path to 4x4 matrix of T_robot_base^arm_base (compose robot poses to robot base)")
    args = parser.parse_args()

    # Load or generate dataset
    if args.data:
        d = np.load(args.data)
        robot_arr = d["robot_poses"]
        cam_arr = d["camera_poses"]
        assert robot_arr.ndim == 3 and robot_arr.shape[1:] == (4, 4)
        assert cam_arr.ndim == 3 and cam_arr.shape[1:] == (4, 4)
        robot_poses = [robot_arr[i] for i in range(robot_arr.shape[0])]
        camera_poses = [cam_arr[i] for i in range(cam_arr.shape[0])]
        print(f"已从 {args.data} 加载数据集，共 {len(robot_poses)} 组")
    elif args.images and args.robot_poses_file and args.camera_calib:
        # Build dataset from images + robot poses
        # 1) images
        if os.path.isdir(args.images):
            img_paths = sorted(glob.glob(os.path.join(args.images, '*')))
        else:
            img_paths = sorted(glob.glob(args.images))
        if not img_paths:
            raise FileNotFoundError(f"未找到图像: {args.images}")
        # 2) robot poses
        robot_poses = load_robot_poses_file(args.robot_poses_file)
        if len(robot_poses) != len(img_paths):
            raise ValueError(f"robot poses 数量({len(robot_poses)})与图像数量({len(img_paths)})不一致。请确保一一对应且顺序匹配。")
        # 3) camera intrinsics
        K, D = load_camera_calib(args.camera_calib)
        # 4) detect and solvePnP
        camera_poses = []
        for img in img_paths:
            T_c_t = solve_pnp_from_image(img, K, D, args.pattern, args.board_rows, args.board_cols, args.square_size, args.charuco_marker)
            camera_poses.append(T_c_t)
        print(f"已从图像与位姿构建数据集，共 {len(camera_poses)} 组")
    else:
        robot_poses, camera_poses = generate_synthetic_dataset(num_poses=args.num_poses, seed=args.seed)
        print(f"生成了 {len(robot_poses)} 组模拟位姿数据 (seed={args.seed})")

    # Optional export of dataset for reproducibility
    if args.export_dataset:
        np.savez(args.export_dataset,
                 robot_poses=np.stack(robot_poses, axis=0),
                 camera_poses=np.stack(camera_poses, axis=0))
        print(f"已导出数据集到: {args.export_dataset}")

    # Frame composition: if robot_poses are in arm_base, optionally compose to robot_base
    reference_base = "robot_base"
    if args.robot_poses_frame == "arm_base":
        reference_base = "arm_base"
        if args.arm_to_robot_base:
            T_rb_ab = load_se3_matrix(args.arm_to_robot_base)  # T_robot_base^arm_base
            robot_poses = [T_rb_ab @ T_ab_g for T_ab_g in robot_poses]
            reference_base = "robot_base"
            print("已将 arm_base 框架下的末端位姿转换到 robot_base 框架")

    def run_mode(mode_name: str) -> np.ndarray:
        print(f"\n=== 标定模式: {mode_name} ===")
        X = hand_eye_calibration(robot_poses, camera_poses, mode=mode_name)
        print("标定结果 X:")
        print(X)
        if mode_name == 'eye_to_hand':
            print(f"参考基座: {reference_base}")
        print(f"平移部分: [{X[0, 3]:.4f}, {X[1, 3]:.4f}, {X[2, 3]:.4f}]")
        return X

    results = {}
    try:
        if args.mode:
            X = run_mode(args.mode)
            results[args.mode] = X
        else:
            X_eih = run_mode('eye_in_hand')
            X_eth = run_mode('eye_to_hand')
            results['eye_in_hand'] = X_eih
            results['eye_to_hand'] = X_eth

        if args.eval:
            # 仅计算 eye-in-hand 残差（需要 X_cam_gripper），若未跑则先跑一次
            if 'eye_in_hand' in results:
                X_cam_gripper = results['eye_in_hand']
            else:
                X_cam_gripper = hand_eye_calibration(robot_poses, camera_poses, mode='eye_in_hand')
            rot_err_deg, trans_err = hand_eye_residual(robot_poses, camera_poses, X_cam_gripper)
            print(f"残差评估：旋转均值 {rot_err_deg:.3f} 度，平移均值 {trans_err:.4f} m")

        if args.out:
            # 若选择了特定模式，写该模式；否则写两者
            payload = {"metadata": {"reference_base": reference_base,
                                     "robot_poses_frame": args.robot_poses_frame,
                                     "composed_with_arm_to_robot_base": bool(args.arm_to_robot_base)}}
            if args.mode:
                payload[args.mode] = to_json_SE3(results[args.mode])
            else:
                payload.update({k: to_json_SE3(v) for k, v in results.items()})
            with open(args.out, 'w') as f:
                json.dump(payload, f, indent=2)
            print(f"已保存标定结果到: {args.out}")

    except Exception as e:
        print(f"标定过程中发生错误: {e}")

    # ==================== 实际应用提示 ====================
    if not args.data:
        print("\n=== 实际应用说明 ===")
        print("1. 从机器人控制器读取末端位姿 (通常是4x4齐次矩阵或位姿向量，注意这是 基座->末端 T_b^g)")
        print("2. 使用 OpenCV solvePnP 从标定板图像解算 T_c^t (标定板到相机):")
        print("   retval, rvec, tvec = cv2.solvePnP(objectPoints, imagePoints, cameraMatrix, distCoeffs)")
        print("   camera_pose = create_homogeneous_matrix(rvec, tvec)  # 即 T_c^t")
        print("3. 收集多组对应的 T_b^g 和 T_c^t（包含多轴较大旋转，>= 8-12 组）")
        print("4. 运行本脚本 --data dataset.npz --mode eye_to_hand --out result.json --eval")
        print("\n若 robot_poses 在 arm_base 框架下：")
        print("- 如无 arm_base->robot_base 变换，脚本将输出 arm_base^T_camera（eye_to_hand）")
        print("- 若提供 --arm-to-robot-base T_robot_base_from_arm_base.npy，将自动转换并输出 robot_base^T_camera")