"""
calc_new_dh.py - 计算新URDF的DH参数和质心坐标转换
"""
import math

def Rx(angle):
    c, s = math.cos(angle), math.sin(angle)
    return [[1, 0, 0], [0, c, -s], [0, s, c]]

def Ry(angle):
    c, s = math.cos(angle), math.sin(angle)
    return [[c, 0, s], [0, 1, 0], [-s, 0, c]]

def Rz(angle):
    c, s = math.cos(angle), math.sin(angle)
    return [[c, -s, 0], [s, c, 0], [0, 0, 1]]

def mat_mul(A, B):
    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):
    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):
    return mat_mul(mat_mul(Rz(yaw), Ry(pitch)), Rx(roll))

def transpose(M):
    return [[M[j][i] for j in range(3)] for i in range(3)]

# ============ NEW URDF (arm.urdf in fix2026224) ============
print("=" * 72)
print("NEW URDF: Link parameters")
print("=" * 72)

new_links = [
    ("Link1", (0.0066971, -0.00030098, 0.044241),  2.3045,  (0, 0, 0)),
    ("Link2", (-3.7501e-8, -0.26586,    0.00044116),0.8903,  (-1.5708, 0, -1.5708)),
    ("Link3", (0.0016734, -0.086963,   -0.025004),  0.72964, (0, 0, -1.5708)),
    ("Link4", (4.9724e-5, -0.0020707,   0.1436),    0.60215, (-1.5708, 0, -1.5708)),
    ("Link5", (-1.7073e-5, 0.058169,    0.0029732), 0.21817, (1.5708, 0, -3.1416)),
    ("Link6", (2.1332e-5,  0.00053897,  0.099935),  0.57513, (1.5708, 0, 3.1416)),
]

results = []
for name, rc_urdf, mass, joint_rpy in new_links:
    roll, pitch, yaw = joint_rpy
    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}  (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(f"  mass = {mass}")
    print()

# ============ Joint origin comparison ============
print("=" * 72)
print("Joint origin comparison: OLD vs NEW")
print("=" * 72)
old_joints = [
    ("j1", (0.18701, -1.0602, 0.042134), (0, 0, 1.5708)),
    ("j2", (-0.0208, 0, 0.085),           (-1.5708, 0, -1.5708)),
    ("j3", (0, -0.404, 0.0563),            (0, 0, -1.5708)),
    ("j4", (0.047775, -0.0905, -0.0355),   (-1.5708, 0, -1.5708)),
    ("j5", (0, 0.023, 0.241),              (1.5708, 0, 3.1416)),
    ("j6", (1.4024e-5, 0.1065, -0.022997), (1.5708, 0, 3.1416)),
]
new_joints = [
    ("j1", (0, 0, 0.0405),                (0, 0, 0)),
    ("j2", (-0.0014, 0, 0.085),            (-1.5708, 0, -1.5708)),
    ("j3", (0, -0.388, 0.002795),          (0, 0, -1.5708)),
    ("j4", (0.047775, -0.0905, -0.001395), (-1.5708, 0, -1.5708)),
    ("j5", (0, -0.004, 0.241),             (1.5708, 0, -3.1416)),
    ("j6", (0, 0.1065, 0.0040025),         (1.5708, 0, 3.1416)),
]

for old, new in zip(old_joints, new_joints):
    name = old[0]
    xyz_old, rpy_old = old[1], old[2]
    xyz_new, rpy_new = new[1], new[2]
    changed = "***CHANGED***" if xyz_old != xyz_new or rpy_old != rpy_new else ""
    print(f"{name}: {changed}")
    print(f"  OLD xyz={xyz_old}  rpy=({math.degrees(rpy_old[0]):.1f}, {math.degrees(rpy_old[1]):.1f}, {math.degrees(rpy_old[2]):.1f})")
    print(f"  NEW xyz={xyz_new}  rpy=({math.degrees(rpy_new[0]):.1f}, {math.degrees(rpy_new[1]):.1f}, {math.degrees(rpy_new[2]):.1f})")
    print()

# ============ Generate C code ============
print("=" * 72)
print("C CODE for arm_main.cpp dh_params")
print("=" * 72)

# DH params extracted from joint origins:
# Standard DH: d comes from z-component of joint origin, a from x/y distance
# These need to be extracted from the URDF joint xyz and rpy carefully
# For the new URDF, the DH parameters:
dh_raw = [
    # d,          a,          alpha,       theta_off,   qmin,   qmax,  effort
    (0.0405,     0.0014,    -math.pi/2,   0.0,         -15.7,  15.7,  60),
    (0.0,        0.388,      0.0,         -math.pi/2,  -1.57,  1.57,  60),
    (0.002795,   0.047775,  -math.pi/2,   -math.pi/2,  0.0,    3.0,   60),
    (0.241,      0.0,        math.pi/2,    math.pi,     0.0,    6.28,  20),
    (0.0,        0.1065,    -math.pi/2,    0.0,         -1.75,  1.75,  20),
    (0.0040025,  0.0,        0.0,          math.pi,     0.0,    6.3,   20),
]

alpha_strs   = ["-M_PI_2", "0.0f", "-M_PI_2", "M_PI_2", "-M_PI_2", "0.0f"]
toffset_strs = ["0.0f", "-M_PI_2", "-M_PI_2", "M_PI", "0.0f", "M_PI"]
comments = [
    "// J1: yaw",
    "// J2: pitch (big arm)",
    "// J3: pitch (forearm)",
    "// J4: roll (wrist)",
    "// J5: pitch (wrist)",
    "// J6: roll (end)",
]

print("static Arm6dof_DHParams_t dh_params[6] = {")
for i, (name, rc_urdf, rc_dh, mass, joint_rpy) in enumerate(results):
    d, a, alpha, toff, qmin, qmax, effort = dh_raw[i]
    rx, ry, rz = rc_dh
    print(f"  {comments[i]}")
    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}}, .kp=1.0f, .kd=0.0f}},")

print("};")