# id 1  120 6500
# id 2  120 9000
# id 3  90 6300


import tkinter as tk
from tkinter import ttk
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.figure import Figure
import numpy as np
from mpl_toolkits.mplot3d import proj3d
from ZDT_Servo import Emm42Driver

class Arm3DVisualizer(tk.Tk):
    def _init_send_queue(self):
        import threading, queue
        self._send_queue = queue.Queue()
        self._send_thread_stop = False
        def send_loop():
            while not self._send_thread_stop:
                try:
                    args = self._send_queue.get(timeout=0.005)
                except Exception:
                    continue
                # 非阻塞发送
                self._send_to_motor_nowait(*args)
        t = threading.Thread(target=send_loop, daemon=True)
        t.start()
        self._send_thread = t

    def _send_to_motor_nowait(self, a1, a2, 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=1000, pulses=pulse, is_absolute=True, acceleration=10)
                except Exception as e:
                    print(f"ID{node_id}异常:{e}")
    def _start_motion_thread(self):
        import threading, time
        self._last_sent_angles = list(self.joint_angles)
        self._motion_thread_stop = False
        self._motion_accum_count = 0
        self._motion_last_time = time.time()
        def motion_loop():
            DEADZONE = 5  # 死区阈值，单位：度
            SEND_INTERVAL = 0.06  # 最小发送间隔，单位：秒
            while not self._motion_thread_stop:
                # 死区判定：每轴变化超过0.5度才算变化
                changed = any(abs(a-b)>DEADZONE for a,b in zip(self.joint_angles, self._last_sent_angles))
                now = time.time()
                # 满足变化且距离上次发送超过60ms才发送
                if changed and (now - self._motion_last_time) > SEND_INTERVAL:
                    self._motion_last_time = now
                    self._last_sent_angles = list(self.joint_angles)
                    self.send_to_motor()
                time.sleep(0.01)
        t = threading.Thread(target=motion_loop, daemon=True)
        t.start()
        self._motion_thread = t

    def _init_key_state(self):
        self._key_state = {'w':False,'s':False,'a':False,'d':False,'r':False,'f':False}
        self._key_step_xy = 10  # 步长
        self._key_step_z = 10
        self.bind_all('<KeyPress>', self._on_key_press)
        self.bind_all('<KeyRelease>', self._on_key_release)
        self._key_move_loop()

    def _on_key_press(self, event):
        k = event.keysym.lower()
        if k in self._key_state:
            self._key_state[k] = True

    def _on_key_release(self, event):
        k = event.keysym.lower()
        if k in self._key_state:
            self._key_state[k] = False

    def _key_move_loop(self):
        # 按住键平滑移动
        if self.pt3d is not None:
            x, y, z = self.pt3d
            moved = False
            if self._key_state['w']:
                y += self._key_step_xy; moved = True
            if self._key_state['s']:
                y -= self._key_step_xy; moved = True
            if self._key_state['a']:
                x -= self._key_step_xy; moved = True
            if self._key_state['d']:
                x += self._key_step_xy; moved = True
            if self._key_state['r']:
                z += self._key_step_z; moved = True
            if self._key_state['f']:
                z -= self._key_step_z; moved = True
            if moved:
                self.inverse_kinematics_and_update(x, y, z)
                self.status_label.config(text=f"目标位置: x={x:.1f}, y={y:.1f}, z={z:.1f}")
        self.after(20, self._key_move_loop)
    def _init_motors(self):
        # 初始化3个Emm42Driver实例
        self.motors = {}
        for i in range(1, 4):
            try:
                self.motors[i] = Emm42Driver(port='/dev/ttyAMA5', device_id=i, baudrate=115200)
            except Exception as e:
                self.motors[i] = None
                print(f"电机{i}初始化失败: {e}")
    def __init__(self):
        super().__init__()
        self._init_send_queue()  # 必须最先初始化
        self.title("三轴机械臂三维可视化与拖拽演示")
        self.geometry("700x600")
        # 机械臂参数
        self.arm_params = {'L1': 185, 'L2': 230, 'L3': 245}
        # 初始角度
        self.joint_angles = [0, 120, -155]
        # 末端位置
        self.pt3d = None
        self.init_ui()
        self._init_motors()
        self.update_arm_plot()
        self._init_key_state()
        self._start_motion_thread()

    def init_ui(self):
        # 主体左右布局
        main_frame = ttk.Frame(self)
        main_frame.pack(fill='both', expand=True)

        # 左侧3D
        left_frame = ttk.Frame(main_frame)
        left_frame.pack(side=tk.LEFT, fill='both', expand=True)
        fig = Figure(figsize=(5, 5), dpi=100)
        self.ax = fig.add_subplot(111, projection='3d')
        self.canvas = FigureCanvasTkAgg(fig, master=left_frame)
        self.canvas.get_tk_widget().pack(fill='both', expand=True)

        # 右侧xy+z
        right_frame = ttk.Frame(main_frame)
        right_frame.pack(side=tk.LEFT, fill='y', padx=5, pady=5)
        # xy平面正方形
        from matplotlib.figure import Figure as Fig2D
        from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg as Canvas2D
        self.xy_fig = Fig2D(figsize=(3,3), dpi=80)
        self.xy_ax = self.xy_fig.add_subplot(111)
        self.xy_canvas = Canvas2D(self.xy_fig, master=right_frame)
        self.xy_canvas.get_tk_widget().pack(side=tk.TOP, fill='both', expand=True)
        self.xy_canvas.get_tk_widget().configure(width=240, height=240)
        self.xy_drag_point = None
        self.xy_dragging = False
        self.xy_canvas.mpl_connect('button_press_event', self._on_xy_press)
        self.xy_canvas.mpl_connect('motion_notify_event', self._on_xy_motion)
        self.xy_canvas.mpl_connect('button_release_event', self._on_xy_release)
        # z轴竖直滑条
        z_frame = ttk.Frame(right_frame)
        z_frame.pack(side=tk.TOP, fill='y', pady=10)
        ttk.Label(z_frame, text='Z轴:').pack(side=tk.TOP)
        zmin = self.arm_params['L1']
        zmax = self.arm_params['L1']+self.arm_params['L2']+self.arm_params['L3']
        self.z_var = tk.DoubleVar()
        self.z_scale = ttk.Scale(z_frame, from_=zmax, to=zmin, orient=tk.VERTICAL, variable=self.z_var, command=self._on_z_change, length=200)
        self.z_scale.pack(side=tk.TOP, fill='y', expand=True)

        # 角度显示
        angle_frame = ttk.Frame(self)
        angle_frame.pack(fill='x', padx=5, pady=5)
        self.angle_vars = []
        for i in range(3):
            ttk.Label(angle_frame, text=f"关节{i+1}角度(°):").pack(side=tk.LEFT)
            var = tk.DoubleVar(value=self.joint_angles[i])
            self.angle_vars.append(var)
            entry = ttk.Entry(angle_frame, textvariable=var, width=6)
            entry.pack(side=tk.LEFT, padx=2)
        ttk.Button(angle_frame, text="刷新姿态", command=self.on_update_angles).pack(side=tk.LEFT, padx=10)
        # 发送到电机按钮
        ttk.Button(angle_frame, text="发送到电机", command=self.send_to_motor).pack(side=tk.LEFT, padx=10)
        self.status_label = ttk.Label(self, text="键盘控制：WSAD前后左右，R/F上下")
        self.status_label.pack(fill='x')
        # 脉冲显示
        self.pulse_label = ttk.Label(self, text="")
        self.pulse_label.pack(fill='x')
        # 初始点居中
        if self.pt3d is None:
            self.pt3d = [0, 0, (zmin+zmax)/2]
        self.z_var.set(self.pt3d[2])
        self._draw_xy_plane()

    def _draw_xy_plane(self):
        # 画xy平面和当前点，保持正方形
        self.xy_ax.clear()
        L = self.arm_params['L2']+self.arm_params['L3']
        self.xy_ax.set_xlim(-L, L)
        self.xy_ax.set_ylim(-L, L)
        self.xy_ax.set_aspect('equal', adjustable='box')
        self.xy_ax.set_xlabel('X (mm)')
        self.xy_ax.set_ylabel('Y (mm)')
        if self.pt3d is not None:
            x, y = self.pt3d[0], self.pt3d[1]
            self.xy_drag_point = self.xy_ax.plot(x, y, 'ro', markersize=10, picker=5)[0]
        self.xy_ax.grid(True)
        self.xy_canvas.draw()

    def _on_xy_press(self, event):
        if event.inaxes != self.xy_ax: return
        if self.pt3d is None: return
        x, y = self.pt3d[0], self.pt3d[1]
        L = self.arm_params['L2']+self.arm_params['L3']
        if abs(event.xdata - x) < L*0.08 and abs(event.ydata - y) < L*0.08:
            self.xy_dragging = True

    def _on_xy_motion(self, event):
        if not self.xy_dragging or event.inaxes != self.xy_ax: return
        z = self.z_var.get() if hasattr(self, 'z_var') else self.pt3d[2]
        self.inverse_kinematics_and_update(event.xdata, event.ydata, z)
        self.z_var.set(z)
        self._draw_xy_plane()

    def _on_xy_release(self, event):
        self.xy_dragging = False

    def _on_z_change(self, val):
        if self.pt3d is not None:
            x, y = self.pt3d[0], self.pt3d[1]
            z = float(val)
            self.inverse_kinematics_and_update(x, y, z)
            self._draw_xy_plane()

    def send_to_motor(self):
        # 只将当前角度放入队列，不等待
        self._send_queue.put(tuple(self.joint_angles))
    def destroy(self):
        # 停止后台线程
        self._motion_thread_stop = True
        self._send_thread_stop = True
        super().destroy()

    def on_update_angles(self):
        try:
            # 限位
            angles = [var.get() for var in self.angle_vars]
            angles[1] = min(120, angles[1])
            angles[2] = max(-155, min(0, angles[2]))
            self.joint_angles = angles
            for i in range(3):
                self.angle_vars[i].set(self.joint_angles[i])
            self.update_arm_plot()
            self.update_pulse_label()
        except Exception as e:
            self.status_label.config(text=f"角度输入错误: {e}")

    def forward_kinematics(self, theta1, theta2, theta3):
        L1, L2, L3 = self.arm_params['L1'], self.arm_params['L2'], self.arm_params['L3']
        theta1 = np.deg2rad(theta1)
        theta2 = np.deg2rad(theta2)
        theta3 = np.deg2rad(theta3)
        x0, y0, z0 = 0, 0, 0
        x1, y1, z1 = 0, 0, L1

        x2 = L2 * np.cos(theta2) * np.cos(theta1)
        y2 = L2 * np.cos(theta2) * np.sin(theta1)
        z2 = L1 + L2 * np.sin(theta2)
        x3 = x2 + L3 * np.cos(theta2 + theta3) * np.cos(theta1)
        y3 = y2 + L3 * np.cos(theta2 + theta3) * np.sin(theta1)
        z3 = z2 + L3 * np.sin(theta2 + theta3)
        return np.array([[x0, y0, z0], [x1, y1, z1], [x2, y2, z2], [x3, y3, z3]])

    def update_arm_plot(self):
        pts = self.forward_kinematics(*self.joint_angles)
        # 保留当前视角
        elev = self.ax.elev
        azim = self.ax.azim
        self.ax.clear()
        self.ax.plot(pts[:,0], pts[:,1], pts[:,2], '-o', linewidth=3, markersize=8, color='b')
        self.ax.scatter(pts[3,0], pts[3,1], pts[3,2], s=120, c='r', picker=5)
        self.ax.set_xlim(-self.arm_params['L2']-self.arm_params['L3'], self.arm_params['L2']+self.arm_params['L3'])
        self.ax.set_ylim(-self.arm_params['L2']-self.arm_params['L3'], self.arm_params['L2']+self.arm_params['L3'])
        self.ax.set_zlim(0, self.arm_params['L1']+self.arm_params['L2']+self.arm_params['L3'])
        self.ax.set_xlabel('X (mm)')
        self.ax.set_ylabel('Y (mm)')
        self.ax.set_zlabel('Z (mm)')
        self.ax.set_title('三轴机械臂三维可视化')
        self.ax.view_init(elev=elev, azim=azim)
        self.canvas.draw()
        self.pt3d = pts[3]
        self.update_pulse_label()

    def update_pulse_label(self):
        # 角度转脉冲
        a1, a2, a3 = self.joint_angles
        # id1: 120度=6500，负角度为正脉冲
        pulse1 = int(6500 * (a1 / 120))
        # id2: 120度=9000，120度为0，脉冲增加角度减小
        pulse2 = int(9000 * (120 - a2) / 120)
        # id3: -155度=0，-65度=6300，脉冲增加角度增加
        # 线性关系: pulse = (a3 + 155) / ( -65 + 155 ) * 6300, 但-65-(-155)=90
        pulse3 = int(6300 * (a3 + 155) / 90)
        self.pulse_label.config(text=f"脉冲值: ID1={pulse1}  ID2={pulse2}  ID3={pulse3}")
        # 自动发送到电机
        # self.send_to_motor()


    # 键盘事件已由_key_move_loop和_key_state管理
    def destroy(self):
        # 停止后台线程
        self._motion_thread_stop = True
        super().destroy()

    def inverse_kinematics_and_update(self, x, y, z):
        L1, L2, L3 = self.arm_params['L1'], self.arm_params['L2'], self.arm_params['L3']
        theta1 = np.arctan2(y, x)
        r = np.hypot(x, y)
        dz = z - L1
        D = (r**2 + dz**2 - L2**2 - L3**2) / (2*L2*L3)
        if np.abs(D) > 1:
            self.status_label.config(text="超出机械臂工作空间")
            return
        theta3 = np.arctan2(-np.sqrt(1-D**2), D)  # 肘下解
        theta2 = np.arctan2(dz, r) - np.arctan2(L3*np.sin(theta3), L2+L3*np.cos(theta3))
        # 限位
        theta2_deg = min(120, np.rad2deg(theta2))
        theta3_deg = max(-155, min(0, np.rad2deg(theta3)))
        self.joint_angles = [np.rad2deg(theta1), theta2_deg, theta3_deg]
        for i in range(3):
            self.angle_vars[i].set(self.joint_angles[i])
        self.update_arm_plot()
        self.update_pulse_label()

if __name__ == "__main__":
    app = Arm3DVisualizer()
    app.mainloop()
