"""
urdf_rc_to_dh.py
================
将 URDF 中各连杆的质心坐标 (rc_urdf) 转换到 DH 连杆坐标系下 (rc_dh)。
纯标准库实现，无需 numpy。

关键修正
--------
URDF中的Link坐标系由parent joint的完整6DOF变换（xyz + rpy）定义。
DH坐标系仅用4个参数（d, a, α, θ）定义。
需要用joint的rpy旋转来转换质心坐标，而不仅仅是theta_offset。

转换关系：
    rc_dh = R_rpy^(-1) * rc_urdf
其中 R_rpy 是URDF joint中定义的旋转矩阵（roll-pitch-yaw顺序）
"""

import math

def Rx(angle):
    """绕x轴旋转的旋转矩阵"""
    c, s = math.cos(angle), math.sin(angle)
    return [[1, 0, 0], [0, c, -s], [0, s, c]]

def Ry(angle):
    """绕y轴旋转的旋转矩阵"""
    c, s = math.cos(angle), math.sin(angle)
    return [[c, 0, s], [0, 1, 0], [-s, 0, c]]

def Rz(angle):
    """绕z轴旋转的旋转矩阵"""
    c, s = math.cos(angle), math.sin(angle)
    return [[c, -s, 0], [s, c, 0], [0, 0, 1]]

def mat_mul(A, B):
    """3x3矩阵乘法"""
    return [[sum(A[i][k] * B[k][j] for k in range(3)) for j in range(3)] for i in range(3)]

def mat_vec_mul(M, v):
    """3x3矩阵乘3x1向量"""
    return tuple(sum(M[i][j] * v[j] for j in range(3)) for i in range(3))

def rpy_to_matrix(roll, pitch, yaw):
    """RPY(roll-pitch-yaw)转旋转矩阵，顺序：Rz(yaw) * Ry(pitch) * Rx(roll)"""
    return mat_mul(mat_mul(Rz(yaw), Ry(pitch)), Rx(roll))

def transpose(M):
    """3x3矩阵转置"""
    return [[M[j][i] for j in range(3)] for i in range(3)]

# ============================================================
# 从 URDF 直接提取的数据
# ============================================================
links_urdf = [
    # name, rc_urdf (x,y,z), mass, joint_rpy (从URDF joint定义)
    ("Link1", ( 0.0070301,  -0.00030378,  0.044333  ),  2.3117,  (0, 0, 1.5708)),        # j1: rpy="0 0 1.5708"
    ("Link2", (-9.2767e-6,  -0.27664,     0.019763  ),  0.9007,  (-1.5708, 0, -1.5708)), # j2: rpy="-1.5708 0 -1.5708"
    ("Link3", ( 0.0016734,  -0.086963,   -0.059109  ),  0.72964, (0, 0, -1.5708)),       # j3: rpy="0 0 -1.5708"
    ("Link4", ( 4.9724e-5,  -0.00207071,  0.143600  ),  0.60215, (-1.5708, 0, -1.5708)), # j4: rpy="-1.5708 0 -1.5708"
    ("Link5", (-1.00614e-5,  0.05816669, -0.02402677),  0.21817, (1.5708, 0, 3.1416)),   # j5: rpy="1.5708 0 3.1416"
    ("Link6", ( 2.13317e-5,  5.38972e-4,  0.099935  ),  0.57513, (1.5708, 0, 3.1416)),   # j6: rpy="1.5708 0 3.1416"
]

print("=" * 72)
print("连杆质心坐标：URDF坐标系 → DH连杆坐标系")
print("=" * 72)
print()

results = []
for name, rc_urdf, mass, joint_rpy in links_urdf:
    roll, pitch, yaw = joint_rpy
    # URDF坐标系 = R_rpy * DH坐标系
    # 所以 rc_dh = R_rpy^T * rc_urdf (逆旋转)
    R = rpy_to_matrix(roll, pitch, yaw)
    R_inv = transpose(R)  # 旋转矩阵的逆 = 转置
    rc_dh = mat_vec_mul(R_inv, rc_urdf)
    results.append((name, rc_urdf, rc_dh, mass, joint_rpy))

    print(f"{name}  (joint rpy = {math.degrees(roll):.1f}°, {math.degrees(pitch):.1f}°, {math.degrees(yaw):.1f}°)")
    print(f"  URDF rc = ({rc_urdf[0]:+.8f}, {rc_urdf[1]:+.8f}, {rc_urdf[2]:+.8f})")
    print(f"  DH   rc = ({rc_dh[0]:+.8f}, {rc_dh[1]:+.8f}, {rc_dh[2]:+.8f})")
    print()

# ============================================================
# 生成可直接复制到 arm_main.cpp 的 C 代码
# ============================================================
print("=" * 72)
print("生成的 C 代码（直接替换 arm_main.cpp 中的 dh_params）")
print("=" * 72)
print()
print("// DH参数 (单位: m, rad, kg)")
print("// rc[3]: 质心在 DH 连杆坐标系下的坐标 (m)，由 URDF inertial/origin 转换")
print("// 转换: rc_dh = R_rpy^(-1) * rc_urdf，其中R_rpy来自URDF joint的rpy")
print("static Arm6dof_DHParams_t dh_params[6] = {")

dh_raw = [
    # d,          a,          alpha,      theta_off,    qmin,   qmax
    ( 0.042134,   0.0208,    -math.pi/2,  math.pi/2,   -15.7,   15.7  ),
    ( 0.0,        0.404,      0.0,        -math.pi/2,  -1.57,    1.57  ),
    ( 0.0563,     0.047775,  -math.pi/2,  -math.pi/2,   0.0,     3.0   ),
    ( 0.241,      0.0,        math.pi/2,   math.pi,    -12.56,  12.56  ),
    ( 0.0,        0.1065,    -math.pi/2,   0.0,         -1.75,   1.75  ),
    ( 0.099935,   0.0,        0.0,         math.pi,    -12.56,  12.56  ),
]

alpha_strs  = ["-M_PI_2", "0.0f", "-M_PI_2", "M_PI_2", "-M_PI_2", "0.0f"]
toffset_strs = ["M_PI_2", "-M_PI_2", "-M_PI_2", "M_PI", "0.0f", "M_PI"]
comments = [
    "// J1: 腰部旋转",
    "// J2: 大臂俯仰",
    "// J3: 小臂",
    "// J4: 腕部旋转",
    "// J5: 腕部俯仰",
    "// J6: 末端Roll",
]

for i, (name, rc_urdf, rc_dh, mass, joint_rpy) in enumerate(results):
    d, a, alpha, toff, qmin, qmax = dh_raw[i]
    rx, ry, rz = rc_dh
    roll, pitch, yaw = joint_rpy
    print(f"  {comments[i]}  URDF rc=({rc_urdf[0]:+.4e}, {rc_urdf[1]:+.4e}, {rc_urdf[2]:+.4e}), rpy=({math.degrees(roll):.0f}°,{math.degrees(pitch):.0f}°,{math.degrees(yaw):.0f}°)")
    print(f"  {{.theta=0.0f, .d={d}f, .a={a}f, .alpha={alpha_strs[i]}, .theta_offset={toffset_strs[i]}, "
          f".qmin={qmin}f, .qmax={qmax}f, .m={mass}f, "
          f".rc={{{rx:.8f}f, {ry:.8f}f, {rz:.8f}f}}}},")

print("};")