import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 定义机械参数
leg_params = {
    'hip': 0.1,    # 髋关节长度
    'thigh': 0.3,   # 大腿长度
    'calf': 0.3     # 小腿长度
}

# 定义关节方向符号 (与C代码中一致)
signs = {
    'left_front': {'hip': 1, 'thigh': 1, 'calf': -1},
    'right_front': {'hip': -1, 'thigh': -1, 'calf': 1},
    'left_rear': {'hip': -1, 'thigh': 1, 'calf': -1},
    'right_rear': {'hip': 1, 'thigh': -1, 'calf': 1}
}

def forward_kinematics(theta, leg_params, sign):
    """正运动学"""
    q1, q2, q3 = theta
    q2 = q2 * sign['thigh']
    q3 = q3 * sign['calf']
    
    l1 = leg_params['hip'] * sign['hip']
    l2 = leg_params['thigh']
    l3 = leg_params['calf']
    
    # 计算足端位置
    x = l2 * np.cos(q2) + l3 * np.cos(q2 + q3)
    y = l1 * np.cos(q1) + np.sin(q1) * (l2 * np.sin(q2) + l3 * np.sin(q2 + q3))
    z = l1 * np.sin(q1) - np.cos(q1) * (l2 * np.sin(q2) + l3 * np.sin(q2 + q3))
    
    foot_pos = np.array([x, -y if sign['hip'] == sign['thigh'] else y, z])
    return foot_pos

def inverse_kinematics(foot_pos, leg_params, sign):
    """逆运动学"""
    px = foot_pos[0]
    py = -foot_pos[1] if sign['hip'] == sign['thigh'] else foot_pos[1]
    pz = foot_pos[2]
    
    b2y = leg_params['hip'] * sign['hip']
    b3z = -leg_params['thigh']
    b4z = -leg_params['calf']
    a = leg_params['hip']
    c = np.sqrt(px**2 + py**2 + pz**2)
    b = np.sqrt(c**2 - a**2)
    
    # 检查可达性
    if np.isnan(b) or b > abs(b3z) + abs(b4z):
        return None
    
    def q1_ik(py, pz, l1):
        L = np.sqrt(py**2 + pz**2 - l1**2)
        return np.arctan2(pz*l1 + py*L, py*l1 - pz*L)
    
    def q3_ik(b3z, b4z, b):
        temp = (b3z**2 + b4z**2 - b**2) / (2.0 * abs(b3z * b4z))
        temp = np.clip(temp, -1.0, 1.0)
        q3 = np.arccos(temp)
        return -(np.pi - q3)
    
    def q2_ik(q1, q3, px, py, pz, b3z, b4z):
        a1 = py * np.sin(q1) - pz * np.cos(q1)
        a2 = px
        m1 = b4z * np.sin(q3)
        m2 = b3z + b4z * np.cos(q3)
        return np.arctan2(m1*a1 + m2*a2, m1*a2 - m2*a1)
    
    q1 = q1_ik(py, pz, b2y)
    q3 = q3_ik(b3z, b4z, b)
    q2 = q2_ik(q1, q3, px, py, pz, b3z, b4z)
    
    return np.array([q1, q2 * sign['thigh'], q3 * sign['calf']])

def test_kinematics(leg_sign, test_name):
    print(f"\nTesting {test_name} leg...")
    
    # 生成随机关节角度
    np.random.seed(42)
    test_angles = np.random.uniform(-np.pi/2, np.pi/2, (10, 3))
    
    errors = []
    
    for angles in test_angles:
        # 正运动学计算足端位置
        foot_pos = forward_kinematics(angles, leg_params, leg_sign)
        
        # 逆运动学计算关节角度
        calculated_angles = inverse_kinematics(foot_pos, leg_params, leg_sign)
        
        if calculated_angles is None:
            print(f"Unreachable position for angles: {angles}")
            continue
        
        # 计算误差
        error = np.max(np.abs(angles - calculated_angles))
        errors.append(error)
        
        print(f"Original angles: {angles}")
        print(f"Calculated angles: {calculated_angles}")
        print(f"Error: {error:.6f} rad\n")
    
    avg_error = np.mean(errors)
    max_error = np.max(errors)
    print(f"Average error: {avg_error:.6f} rad")
    print(f"Maximum error: {max_error:.6f} rad")
    
    return avg_error, max_error

# 测试所有四种腿型
results = {}
for leg_name, leg_sign in signs.items():
    avg_err, max_err = test_kinematics(leg_sign, leg_name)
    results[leg_name] = (avg_err, max_err)

# 打印汇总结果
print("\nSummary of results:")
for leg_name, (avg_err, max_err) in results.items():
    print(f"{leg_name}: avg error={avg_err:.6f} rad, max error={max_err:.6f} rad")

# 可视化一个例子
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# 选择左前腿作为示例
leg_sign = signs['left_front']
angles = np.array([0.5, -0.3, 0.7])  # 髋关节、大腿、小腿角度

# 计算正运动学
foot_pos = forward_kinematics(angles, leg_params, leg_sign)

# 绘制
l1 = leg_params['hip'] * leg_sign['hip']
l2 = leg_params['thigh']
l3 = leg_params['calf']

q1, q2, q3 = angles
q2 = q2 * leg_sign['thigh']
q3 = q3 * leg_sign['calf']

# 计算各关节位置
hip_pos = np.array([0, 0, 0])
knee_pos = np.array([
    l2 * np.cos(q2),
    l1 * np.cos(q1) + np.sin(q1) * l2 * np.sin(q2),
    l1 * np.sin(q1) - np.cos(q1) * l2 * np.sin(q2)
])
foot_pos = np.array([
    l2 * np.cos(q2) + l3 * np.cos(q2 + q3),
    l1 * np.cos(q1) + np.sin(q1) * (l2 * np.sin(q2) + l3 * np.sin(q2 + q3)),
    l1 * np.sin(q1) - np.cos(q1) * (l2 * np.sin(q2) + l3 * np.sin(q2 + q3))
])

# 调整y坐标符号
if leg_sign['hip'] != leg_sign['thigh']:
    knee_pos[1] = -knee_pos[1]
    foot_pos[1] = -foot_pos[1]

# 绘制连杆
ax.plot([hip_pos[0], knee_pos[0]], [hip_pos[1], knee_pos[1]], [hip_pos[2], knee_pos[2]], 'b-', linewidth=3)
ax.plot([knee_pos[0], foot_pos[0]], [knee_pos[1], foot_pos[1]], [knee_pos[2], foot_pos[2]], 'r-', linewidth=3)

# 绘制关节
ax.scatter(hip_pos[0], hip_pos[1], hip_pos[2], c='k', s=100, label='Hip')
ax.scatter(knee_pos[0], knee_pos[1], knee_pos[2], c='k', s=100, label='Knee')
ax.scatter(foot_pos[0], foot_pos[1], foot_pos[2], c='g', s=100, label='Foot')

ax.set_xlabel('X (m)')
ax.set_ylabel('Y (m)')
ax.set_zlabel('Z (m)')
ax.set_title('Leg Kinematics Example (Left Front Leg)')
ax.legend()
plt.tight_layout()
plt.show()