# local_controller.py（运行在本地控制机械臂）
REMOTE_HOST = "10.103.69.239"  # SSH 中配置的 Host（~/.ssh/config）
# import math3d
import socket
import numpy as np
import urx
import time

import urx
import math
import numpy as np

REMOTE_HOST = "localhost"
REMOTE_PORT = 9999
UR_ROBOT_IP = "192.168.1.14"

def pose_to_transform(x, y, z, rx, ry, rz):
    """
    rx, ry, rz 是绕XYZ轴的旋转角，单位弧度
    返回 4x4变换矩阵
    """
    # 旋转矩阵绕X轴
    Rx = np.array([
        [1, 0, 0],
        [0, math.cos(rx), -math.sin(rx)],
        [0, math.sin(rx), math.cos(rx)]
    ])
    # 旋转矩阵绕Y轴
    Ry = np.array([
        [math.cos(ry), 0, math.sin(ry)],
        [0, 1, 0],
        [-math.sin(ry), 0, math.cos(ry)]
    ])
    # 旋转矩阵绕Z轴
    Rz = np.array([
        [math.cos(rz), -math.sin(rz), 0],
        [math.sin(rz), math.cos(rz), 0],
        [0, 0, 1]
    ])

    R = Rz @ Ry @ Rx  # 旋转顺序Z->Y->X

    T = np.eye(4)
    T[:3, :3] = R
    T[0, 3] = x
    T[1, 3] = y
    T[2, 3] = z
    return T

def get_action_from_remote(instruction: str) -> np.ndarray:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.connect((REMOTE_HOST, REMOTE_PORT))
        sock.sendall(instruction.encode())
        data = sock.recv(4096)
        action_str = data.decode()
        action = np.array([float(x) for x in action_str.split(",") if x.strip() != ""])

        return action[:-1]

class MyPose:
    def __init__(self, x, y, z, rx, ry, rz):
        self.position = (x, y, z)  # 位置 (米)
        self.orientation = (rx, ry, rz)  # 姿态 (弧度)
    
    def to_tuple(self):
        # 返回 urx 可识别的位姿元组格式：((x,y,z), (rx, ry, rz))
        return (self.position, self.orientation)
    
from datetime import datetime
def main():
    ur = urx.Robot(UR_ROBOT_IP)
    instruction = input("请输入任务描述（如：adjust the end-effector camera to frame a person who looks distressed）：").strip()
    try:
        while True:
            #Adjust the end-effector camera so that the person who looks distressed is centered in the frame and the camera faces them directly.
            print("current ur:",ur.getj())
            if not instruction:
                continue



            print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] 发送任务到远程模型服务器...")
            action = get_action_from_remote(instruction)
            print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] 接收到动作: {action}")

            if len(action) != 6:
                print("⚠️ 动作长度错误，跳过执行")
                continue
            #ur.movej(action.tolist(), acc=0.05, vel=0.05)
            actin_tuple = tuple(action.tolist())
            
            x, y, z, rx, ry, rz = actin_tuple 
            x = x   # 转换为毫米
            y = y  
            z = z 
            rx =rx
            ry =  ry
            rz =rz
            print("actin_tuple=",(x, y, z, rx, ry, rz))
            ur.movej((x, y, z, rx, ry, rz), 0.05, 0.01,relative=False,wait=False)

    except KeyboardInterrupt:
        print("终止控制")
    finally:
        ur.close()

if __name__ == "__main__":
    main()