# -*- coding: utf-8 -*-
"""
UR机器人连接测试脚本
功能：连接指定IP的UR机器人，获取并打印当前TCP位姿
日期：2025.4.24
"""

# 导入所需库
import rtde_control
import rtde_receive
import time
import numpy as np


def connect_robot(robot_ip):
    """
    建立与UR机器人的RTDE连接
    
    参数:
        robot_ip (str): 机器人的IP地址
    
    返回:
        tuple: (rtde_control实例, rtde_receive实例)，连接失败则返回(None, None)
    """
    try:
        # 尝试创建控制和接收接口的实例
        print(f"conecting {robot_ip}...")
        rtde_c = rtde_control.RTDEControlInterface(robot_ip)
        print(rtde_c)
        rtde_r = rtde_receive.RTDEReceiveInterface(robot_ip)
        print(rtde_r)
        print(f"✅ 成功连接到机器人: {robot_ip}")
        return rtde_c, rtde_r
    except Exception as e:
        print(f"❌ 机器人连接失败: {e}")
        return None, None


def test_robot_connection(robot_ip):
    """
    测试机器人连接并获取TCP位姿
    
    参数:
        robot_ip (str): 机器人的IP地址
    """
    print("starting connection test...")
    # 连接机器人
    rtde_c, rtde_r = connect_robot(robot_ip)
    print("waiting for the connection to be established...")
    # 检查连接状态
    if rtde_c is None or rtde_r is None:
        print("🔴 连接失败，测试终止")
        return
    
    try:
        # 获取机器人的当前TCP位置（x,y,z,rx,ry,rz）
        tcp_pose = rtde_r.getActualTCPPose()
        print(f"🟢 当前TCP位姿: {np.round(tcp_pose, 4)}")  # 保留4位小数，更易读
    except Exception as e:
        print(f"🔴 获取TCP位姿时出错: {e}")
    finally:
        # 确保在退出前停止RTDE控制，释放资源
        if rtde_c is not None:
            rtde_c.stopScript()
            print("🔵 已停止RTDE脚本，释放连接资源")
        print("📋 测试完成")
    return tcp_pose


if __name__ == "__main__":
    # 机器人IP地址（请替换为实际的机器人IP）
    ROBOT_IP = "192.168.1.14"
    
    # 执行连接测试
    tcp_pose = test_robot_connection(ROBOT_IP)
    print("tcp_pose=", tcp_pose)