import rclpy
from rclpy.node import Node
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from osc4py3.as_eventloop import *
from osc4py3 import oscmethod as osm
import roboticstoolbox as rtb
import spatialmath as sm
import xml.etree.ElementTree as ET
import time

class ScaledJointTrajectoryPublisher(Node):
    """Node to publish joint trajectories based on OSC messages."""
    def __init__(self, robot, joint_names):
        super().__init__('scaled_joint_trajectory_publisher')
        
        # ROS2 Publisher
        self.publisher = self.create_publisher(
            JointTrajectory, 
            '/scaled_joint_trajectory_controller/joint_trajectory', 
            10
        )

        # Store received joint positions
        self.joint_positions = [0.0] * len(joint_names)
        self.joint_names = joint_names
        self.robot = robot

        osc_startup()
        osc_udp_server("0.0.0.0", 8000, "osc_server")
        print("Server started on 0.0.0.0:8000 \nready to receive messages in the following format: /tcp_coordinates [x, y, z, roll, pitch, yaw] optional: duration as last argument: /tcp_coordinates [x, y, z, roll, pitch, yaw, duration]")
        # Register OSC handler
        osc_method("/tcp_coordinates", self.tcp_coordinates_handler, argscheme=osm.OSCARG_DATAUNPACK)

    def tcp_coordinates_handler(self, *args):
        """Handles incoming OSC messages for tcp position."""
        time1 = time.time()
        if len(args) == len(self.joint_positions):
            x, y, z, roll, pitch, yaw = args
            duration = 4.0  # Default duration
        elif len(args) == len(self.joint_positions) + 1:
            x, y, z, roll, pitch, yaw, duration = args
        else:
            print("Invalid number of arguments")
            return

        # Create the desired end-effector pose
        Tep = sm.SE3(x, y, z) * sm.SE3.RPY([roll, pitch, yaw], order='xyz')

        # Compute the inverse kinematics to get the joint angles
        #time1 = time.time()
        #sol = self.robot.ikine_LM(Tep)
        #print(f"Time taken for ERobot: {time.time() - time1}")
        #time1 = time.time()
        #sol = self.robot.ikine_LM(Tep, q0=self.joint_positions)
        #print(f"Time taken for ERobot with initial guess: {time.time() - time1}")
        #time1 = time.time()
        #sol = self.robot.ets().ikine_LM(Tep)
        #print(f"Time taken for ETS: {time.time() - time1}")
        #time1 = time.time()
        sol = self.robot.ik_LM(Tep, q0=self.joint_positions)
        #print(f"Time taken for ETS with initial guess: {time.time() - time1}")
        if sol[1]:
            joint_positions = list(sol[0])
            self.send_trajectory(joint_positions, duration)
            #print(f"Computed joint positions: {joint_positions}")
        else:
            print("Inverse kinematics failed")
        print(f"Frequency: {1/(time.time() - time1)} Hz")

    def send_trajectory(self, joint_positions, duration=4.0):
        """Publish a joint trajectory command to move the robot."""
        msg = JointTrajectory()
        msg.joint_names = self.joint_names
        point = JointTrajectoryPoint()
        point.positions = joint_positions  # Updated joint positions
        point.time_from_start.sec = int(duration)
        point.time_from_start.nanosec = int((duration - int(duration)) * 1e9)
        
        msg.points.append(point)
        self.publisher.publish(msg)
        print(f"Updated joint positions: {joint_positions}")

def main():
    """Main function to get joint names and start the ROS 2 & OSC system."""

    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']
    robot = rtb.ERobot.URDF('/BA/robot.urdf')

    rclpy.init()

    node = ScaledJointTrajectoryPublisher(robot.ets(), joint_names)
    
    # Run both ROS 2 and OSC Server together
    try:
        while rclpy.ok():
            osc_process()  # Handle one OSC request at a time
            rclpy.spin_once(node, timeout_sec=0.1)  # Process ROS callbacks
    except KeyboardInterrupt:
        print("")
    finally:
        node.destroy_node()
        rclpy.shutdown()
        osc_terminate()

if __name__ == '__main__':
    main()