#!/usr/bin/env python3
"""相机坐标系到机械臂基坐标系的直接变换工具。

使用场景：
1. 已经通过其他方式获得了相机坐标系坐标
2. 测试和验证坐标变换算法
3. 调试机械臂控制逻辑

输入：
    - 相机坐标系坐标 (xc, yc, zc)，单位 mm
    - 当前 TCP 位姿 (x, y, z, phi)

输出：
    - 基坐标系坐标 (xb, yb, zb)，单位 mm
"""

import argparse
import math
import sys
import json
from pathlib import Path
from typing import Tuple

import numpy as np


def euler_to_rotation_matrix(roll_deg: float, pitch_deg: float, yaw_deg: float) -> np.ndarray:
    """欧拉角转旋转矩阵（ZYX顺序）"""
    roll = math.radians(roll_deg)
    pitch = math.radians(pitch_deg)
    yaw = math.radians(yaw_deg)

    Rx = np.array([
        [1, 0, 0],
        [0, math.cos(roll), -math.sin(roll)],
        [0, math.sin(roll), math.cos(roll)]
    ])

    Ry = np.array([
        [math.cos(pitch), 0, math.sin(pitch)],
        [0, 1, 0],
        [-math.sin(pitch), 0, math.cos(pitch)]
    ])

    Rz = np.array([
        [math.cos(yaw), -math.sin(yaw), 0],
        [math.sin(yaw), math.cos(yaw), 0],
        [0, 0, 1]
    ])

    return Rz @ Ry @ Rx


def camera_to_tcp(
    xc: float, yc: float, zc: float,
    tx: float = 0.0, ty: float = 0.0, tz: float = 0.0,
    roll: float = 0.0, pitch: float = 0.0, yaw: float = 0.0
) -> Tuple[float, float, float]:
    """相机坐标系 → TCP 坐标系

    Args:
        xc, yc, zc: 相机坐标系坐标 (mm)
        tx, ty, tz: 相机到 TCP 的平移 (mm)
        roll, pitch, yaw: 相机到 TCP 的旋转 (度)

    Returns:
        (xt, yt, zt): TCP 坐标系坐标 (mm)
    """
    R = euler_to_rotation_matrix(roll, pitch, yaw)
    T = np.array([tx, ty, tz])
    P_cam = np.array([xc, yc, zc])
    P_tcp = R @ P_cam + T
    return float(P_tcp[0]), float(P_tcp[1]), float(P_tcp[2])


def tcp_to_base(
    xt: float, yt: float, zt: float,
    tcp_x: float, tcp_y: float, tcp_z: float, tcp_phi_deg: float
) -> Tuple[float, float, float]:
    """TCP 坐标系 → 机械臂基坐标系（水平相机版本）

    坐标映射：
        - TCP X (相机右) → 基坐标系 垂直于 phi 方向
        - TCP Y (相机下) → 基坐标系 -Z 方向（向下）
        - TCP Z (相机前) → 基坐标系 phi 方向

    Args:
        xt, yt, zt: TCP 坐标系坐标 (mm)
        tcp_x, tcp_y, tcp_z: TCP 当前位置 (mm)
        tcp_phi_deg: TCP 当前朝向角 (度)

    Returns:
        (xb, yb, zb): 基坐标系坐标 (mm)
    """
    phi = math.radians(tcp_phi_deg)

    # 旋转矩阵：TCP → 基坐标系
    R_tcp_to_base = np.array([
        [-math.sin(phi), 0, math.cos(phi)],
        [ math.cos(phi), 0, math.sin(phi)],
        [0,             -1, 0]
    ])

    P_tcp = np.array([xt, yt, zt])
    P_base_relative = R_tcp_to_base @ P_tcp

    P_base = P_base_relative + np.array([tcp_x, tcp_y, tcp_z])

    return float(P_base[0]), float(P_base[1]), float(P_base[2])


def camera_to_base(
    xc: float, yc: float, zc: float,
    tcp_x: float, tcp_y: float, tcp_z: float, tcp_phi_deg: float,
    cam_tx: float = 0.0, cam_ty: float = 0.0, cam_tz: float = 0.0,
    cam_roll: float = 0.0, cam_pitch: float = 0.0, cam_yaw: float = 0.0
) -> Tuple[float, float, float]:
    """完整变换：相机坐标系 → 基坐标系

    Args:
        xc, yc, zc: 相机坐标系坐标 (mm)
        tcp_x, tcp_y, tcp_z, tcp_phi_deg: TCP 当前位姿
        cam_tx, cam_ty, cam_tz: 相机到 TCP 的平移
        cam_roll, cam_pitch, cam_yaw: 相机到 TCP 的旋转

    Returns:
        (xb, yb, zb): 基坐标系坐标 (mm)
    """
    # 步骤 1: 相机 → TCP
    xt, yt, zt = camera_to_tcp(xc, yc, zc, cam_tx, cam_ty, cam_tz,
                                cam_roll, cam_pitch, cam_yaw)

    # 步骤 2: TCP → 基坐标系
    xb, yb, zb = tcp_to_base(xt, yt, zt, tcp_x, tcp_y, tcp_z, tcp_phi_deg)

    return xb, yb, zb


def load_tcp_pose_from_state() -> Tuple[float, float, float, float]:
    """从状态缓存文件读取当前 TCP 位姿"""
    state_file = Path("tools/.udp_control_state.json")

    if not state_file.exists():
        raise FileNotFoundError(
            f"状态文件不存在: {state_file}\n"
            "请先运行一次 udp_control.py 以初始化状态"
        )

    state = json.loads(state_file.read_text())

    # 导入必要的模块计算正运动学
    sys.path.insert(0, str(Path(__file__).parent))
    from udp_control import (
        forward_kinematics,
        ArmGeometry,
        ArmMathState,
        DEFAULT_L1, DEFAULT_L2, DEFAULT_X4, DEFAULT_Z4,
        DEFAULT_ZERO_J2, DEFAULT_ZERO_J3, DEFAULT_ZERO_J4
    )

    geometry = ArmGeometry(l1=DEFAULT_L1, l2=DEFAULT_L2, x4=DEFAULT_X4, z4=DEFAULT_Z4)
    math_state = ArmMathState(
        d1=state["height"],
        theta2_deg=state["j2"] - DEFAULT_ZERO_J2,
        theta3_deg=state["j3"] - DEFAULT_ZERO_J3,
        theta4_deg=state["j4"] - DEFAULT_ZERO_J4,
    )

    pose = forward_kinematics(geometry, math_state)
    return pose.x, pose.y, pose.z, pose.phi_deg


def main() -> int:
    parser = argparse.ArgumentParser(
        description="相机坐标系到机械臂基坐标系的直接变换",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:

  1. 基本用法（手动指定 TCP 位姿）:
     python %(prog)s --xc 10 --yc 5 --zc 250 \\
         --tcp-x 150 --tcp-y 50 --tcp-z -100 --tcp-phi 45

  2. 自动读取 TCP 位姿（从状态缓存）:
     python %(prog)s --xc 10 --yc 5 --zc 250 --auto-tcp

  3. 包含相机到 TCP 的变换:
     python %(prog)s --xc 10 --yc 5 --zc 250 --auto-tcp \\
         --cam-tx 20 --cam-ty 10 --cam-tz 5

相机坐标系定义（水平安装）:
  Xc: 向右
  Yc: 向下
  Zc: 向前（水平，光轴方向）

机械臂基坐标系定义:
  X, Y: 水平面
  Z: 向上
        """
    )

    # 相机坐标系输入
    parser.add_argument("--xc", type=float, required=True,
                       help="相机坐标系 X 坐标（向右，mm）")
    parser.add_argument("--yc", type=float, required=True,
                       help="相机坐标系 Y 坐标（向下，mm）")
    parser.add_argument("--zc", type=float, required=True,
                       help="相机坐标系 Z 坐标（向前，mm）")

    # TCP 位姿
    tcp_group = parser.add_mutually_exclusive_group(required=True)
    tcp_group.add_argument("--auto-tcp", action="store_true",
                          help="自动从状态缓存读取当前 TCP 位姿")
    tcp_group.add_argument("--tcp-manual", action="store_true",
                          help="手动指定 TCP 位姿（需要 --tcp-x/y/z/phi）")

    parser.add_argument("--tcp-x", type=float,
                       help="TCP X 坐标（mm）")
    parser.add_argument("--tcp-y", type=float,
                       help="TCP Y 坐标（mm）")
    parser.add_argument("--tcp-z", type=float,
                       help="TCP Z 坐标（mm）")
    parser.add_argument("--tcp-phi", type=float,
                       help="TCP 朝向角（度）")

    # 相机到 TCP 的变换（可选）
    parser.add_argument("--cam-tx", type=float, default=0.0,
                       help="相机到 TCP 平移 X（mm，默认 0）")
    parser.add_argument("--cam-ty", type=float, default=0.0,
                       help="相机到 TCP 平移 Y（mm，默认 0）")
    parser.add_argument("--cam-tz", type=float, default=0.0,
                       help="相机到 TCP 平移 Z（mm，默认 0）")
    parser.add_argument("--cam-roll", type=float, default=0.0,
                       help="相机到 TCP 旋转 roll（度，默认 0）")
    parser.add_argument("--cam-pitch", type=float, default=0.0,
                       help="相机到 TCP 旋转 pitch（度，默认 0）")
    parser.add_argument("--cam-yaw", type=float, default=0.0,
                       help="相机到 TCP 旋转 yaw（度，默认 0）")

    # 输出格式
    parser.add_argument("--json", action="store_true",
                       help="以 JSON 格式输出")
    parser.add_argument("--quiet", action="store_true",
                       help="只输出坐标，不输出说明信息")

    args = parser.parse_args()

    # 获取 TCP 位姿
    if args.auto_tcp:
        try:
            tcp_x, tcp_y, tcp_z, tcp_phi = load_tcp_pose_from_state()
            if not args.quiet:
                print(f"从状态缓存读取 TCP 位姿: "
                      f"({tcp_x:.3f}, {tcp_y:.3f}, {tcp_z:.3f}), phi={tcp_phi:.3f}°")
        except Exception as e:
            print(f"错误: 无法读取 TCP 位姿: {e}", file=sys.stderr)
            return 1
    else:
        if not all([args.tcp_x is not None, args.tcp_y is not None,
                   args.tcp_z is not None, args.tcp_phi is not None]):
            print("错误: 使用 --tcp-manual 时必须指定 --tcp-x, --tcp-y, --tcp-z, --tcp-phi",
                  file=sys.stderr)
            return 1
        tcp_x, tcp_y, tcp_z, tcp_phi = args.tcp_x, args.tcp_y, args.tcp_z, args.tcp_phi

    # 执行变换
    try:
        xb, yb, zb = camera_to_base(
            args.xc, args.yc, args.zc,
            tcp_x, tcp_y, tcp_z, tcp_phi,
            args.cam_tx, args.cam_ty, args.cam_tz,
            args.cam_roll, args.cam_pitch, args.cam_yaw
        )

        # 输出结果
        if args.json:
            result = {
                "input": {
                    "camera": {"x": args.xc, "y": args.yc, "z": args.zc},
                    "tcp": {"x": tcp_x, "y": tcp_y, "z": tcp_z, "phi": tcp_phi}
                },
                "output": {
                    "base": {"x": xb, "y": yb, "z": zb}
                }
            }
            print(json.dumps(result, indent=2))
        elif args.quiet:
            print(f"{xb:.3f} {yb:.3f} {zb:.3f}")
        else:
            print("=" * 70)
            print("相机坐标系到基坐标系变换")
            print("=" * 70)
            print(f"\n输入（相机坐标系）:")
            print(f"  X_cam = {args.xc:8.3f} mm  (向右)")
            print(f"  Y_cam = {args.yc:8.3f} mm  (向下)")
            print(f"  Z_cam = {args.zc:8.3f} mm  (向前，光轴)")

            print(f"\nTCP 位姿:")
            print(f"  X_tcp = {tcp_x:8.3f} mm")
            print(f"  Y_tcp = {tcp_y:8.3f} mm")
            print(f"  Z_tcp = {tcp_z:8.3f} mm")
            print(f"  φ_tcp = {tcp_phi:8.3f} °")

            if any([args.cam_tx, args.cam_ty, args.cam_tz,
                   args.cam_roll, args.cam_pitch, args.cam_yaw]):
                print(f"\n相机到 TCP 变换:")
                print(f"  平移: ({args.cam_tx:.3f}, {args.cam_ty:.3f}, {args.cam_tz:.3f}) mm")
                print(f"  旋转: roll={args.cam_roll:.3f}°, pitch={args.cam_pitch:.3f}°, yaw={args.cam_yaw:.3f}°")

            print(f"\n输出（机械臂基坐标系）:")
            print(f"  X_base = {xb:8.3f} mm")
            print(f"  Y_base = {yb:8.3f} mm")
            print(f"  Z_base = {zb:8.3f} mm")
            print("=" * 70)

            print(f"\n使用此坐标控制机械臂:")
            print(f"  python tools/udp_control.py pose \\")
            print(f"    --x {xb:.1f} --y {yb:.1f} --z {zb:.1f} \\")
            print(f"    --phi {tcp_phi:.1f}")

        return 0

    except Exception as e:
        print(f"错误: {e}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())
