import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
from osc4py3.as_eventloop import *
from osc4py3 import oscbuildparse
import roboticstoolbox as rtb
import xml.etree.ElementTree as ET
import numpy as np
from scipy.spatial.transform import Rotation as R

class JointStateOSC(Node):
    def __init__(self, robot, joint_names):
        super().__init__('joint_states_osc')

        self.joint_names_urdf = joint_names
        self.robot = robot

        # Create a ROS 2 subscriber to /joint_states topic
        self.subscription = self.create_subscription(
            JointState,
            '/joint_states',
            self.joint_states_callback,
            1  # Queue size
        )

        # Open Sound Control (OSC) Client settings
        self.osc_ip = "127.0.0.1"   # Replace with the target IP
        self.osc_port = 8000        # Replace with the target port

        # Start the OSC system
        osc_startup()

        # Make client channels to send packets
        osc_udp_client(self.osc_ip, self.osc_port, "osc_client")

    def joint_states_callback(self, msg):
        """Callback function to handle incoming joint states."""
        header = msg.header
        joint_names = msg.name
        joint_positions = msg.position
        joint_positions = [joint_positions[joint_names.index(joint)] for joint in self.joint_names_urdf]
        tcp_pos = self.robot.fkine(joint_positions) #, end='ft_frame')
        tcp_xyz = tcp_pos.t
        tcp_rot = tcp_pos.R
        rotation_vector = R.from_matrix(tcp_rot).as_rotvec()
        translation = oscbuildparse.OSCMessage("/tcp_position_t", None, tcp_xyz.tolist())
        osc_send(translation, "osc_client")
        rotation = oscbuildparse.OSCMessage("/tcp_position_R", None, rotation_vector.tolist())
        osc_send(rotation, "osc_client")
        osc_process()
        #print(f"Published TCP position: {tcp_pos}")
        


def main():
    rclpy.init()
    robot = rtb.ERobot.URDF('/BA/robot.urdf')
    tree = ET.parse('/BA/robot.urdf')
    root = tree.getroot()
    joint_names = [joint.get('name') for joint in root.findall('joint') if joint.get('type') == 'revolute' or joint.get('type') == 'continuous' or joint.get('type') == 'prismatic']

    node = JointStateOSC(robot, joint_names)

    print(f"Publishing TCP coordinates to OSC on {node.osc_ip}:{node.osc_port}...")
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        print("shutting down...")
    finally:
        node.destroy_node()
        rclpy.shutdown()
        osc_terminate()

if __name__ == '__main__':
    main()