from ZDT_Servo import Emm42Driver

class Arm:
    def __init__(self, port='/dev/ttyAMA5', baudrate=115200):
        # 初始化3个电机
        self.motors = {}
        for i in range(1, 4):
            try:
                self.motors[i] = Emm42Driver(port=port, device_id=i, baudrate=baudrate)
            except Exception as e:
                self.motors[i] = None
                print(f"电机{i}初始化失败: {e}")

    def set_angles(self, a1, a2, a3, speed_rpm=1000, acceleration=20):
        # 限位
        a1 = max(-120, min(120, a1))
        a2 = min(120, a2)
        a3 = max(-155, min(0, a3))
        # 角度转脉冲，id1允许负脉冲
        raw1 = 6500 * (a1 / 120)
        if raw1 >= 0:
            pulse1 = int(raw1)
            dir1 = 1
        else:
            pulse1 = int(-raw1)
            dir1 = 0
        pulse2 = int(9000 * (120 - a2) / 120)
        pulse3 = int(6300 * (a3 + 155) / 90)
        for node_id, pulse, direction in zip([1, 2, 3], [pulse1, pulse2, pulse3], [dir1, 1, 1]):
            m = self.motors.get(node_id)
            if m:
                try:
                    m.set_position_mode(direction=direction, speed_rpm=speed_rpm, pulses=pulse, is_absolute=True, acceleration=acceleration)
                except Exception as e:
                    print(f"ID{node_id}异常:{e}")

if __name__ == "__main__":
    import tkinter as tk
    arm = Arm()
    # 初始角度
    angles = [0, 120, -155]
    step = [2, 2, 2]  # 步长
    def send():
        arm.set_angles(*angles)
    def on_key(event):
        k = event.keysym.lower()
        if k == 'q':
            angles[0] += step[0]
        elif k == 'a':
            angles[0] -= step[0]
        elif k == 'w':
            angles[1] += step[1]
        elif k == 's':
            angles[1] -= step[1]
        elif k == 'e':
            angles[2] += step[2]
        elif k == 'd':
            angles[2] -= step[2]
        send()
        print(f"当前角度: {angles}")
    root = tk.Tk()
    root.title("三轴机械臂键盘控制示例")
    root.geometry("300x100")
    label = tk.Label(root, text="Q/A:轴1  W/S:轴2  E/D:轴3")
    label.pack()
    root.bind('<KeyPress>', on_key)
    send()
    root.mainloop()
