#!/usr/bin/env python3
import sys
import moveit_commander
import rospy

"""根据轨迹点的六轴角度值，计算末端执行器的位姿"""
"""1.roslaunch lstrobot_moveit_config demo.launch"""
"""2.rosrun lstrobot_planning joint2ee.py"""

class MoveitPlanner:
    def __init__(self, group_name):
        # 初始化 moveit_commander
        moveit_commander.roscpp_initialize(sys.argv)
        rospy.init_node('1111_node', anonymous=True)
        self.robot = moveit_commander.RobotCommander()
        self.group = moveit_commander.MoveGroupCommander(group_name)
        self.planning_scene_interface = moveit_commander.PlanningSceneInterface()
        self.end_effector_link = self.group.get_end_effector_link()

    def get_robot_state(self):
        # 获取当前机械臂状态
        return self.robot.get_current_state()

    def get_end_effector_pose(self):
        # 获取末端执行器的位姿
        return self.group.get_current_pose(self.end_effector_link)

        
#从文件中读取关节角度值
def read_joint_value(input_filename):
    joint_value_list = []
    with open(input_filename, 'r') as file:
        for line in file:
            #将每行的关节值转换为浮点数并去掉换行符
            joint_value=[float(val) for val in line.strip().split()]
            joint_value_list.append(joint_value)
            
    return joint_value_list

#写入末端执行器位姿到目标文件
def write_ee_pose(joint_value_list, planner, output_filename):
    with open(output_filename, 'w') as file:
       for idx,joint_value in enumerate(joint_value_list):
           #设置关节值为当前轨迹点的关节位置
           planner.group.set_joint_value_target(joint_value)
           
           plan=planner.group.plan()
           
           if plan:
                planner.group.execute(plan[1], wait=True)
                
                fk_result = planner.get_end_effector_pose()
                position=fk_result.pose.position
                orientation=fk_result.pose.orientation
           
                #格式化位姿数据
                pose_data = f"Point {idx + 1}:\n"
                pose_data += f"Position: x={position.x}, y={position.y}, z={position.z}\n"
                pose_data += f"Orientation (Quaternion): x={orientation.x}, y={orientation.y}, z={orientation.z}, w={orientation.w}\n"
                pose_data += "-" * 50 + "\n"
                
                file.write(pose_data)
                
    print(f"End effector pose saved to {output_filename}.")
    
def main():
    planner = MoveitPlanner("manipulator")
    
    input_filename = '/home/tong/moveir_ws/data/123/outtt.txt'
    joint_value_list = read_joint_value(input_filename)
    
    output_filename = '/home/tong/moveir_ws/data/123/ee_pose.txt'
    write_ee_pose(joint_value_list, planner, output_filename)
    
if __name__ == '__main__':
    main()