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 xml.etree.ElementTree as ET
import roboticstoolbox as rtb
import spatialmath as sm
import numpy as np
import time

class ScaledJointTrajectoryPublisher(Node):
    """Node to publish joint trajectories based on OSC messages."""
    def __init__(self, joint_names, robot):
        super().__init__('scaled_joint_trajectory_publisher')
        
        self.robot = robot

        # ROS2 Publisher
        self.publisher = self.create_publisher(
            JointTrajectory, 
            '/scaled_joint_trajectory_controller/joint_trajectory', 
            10
        )

        # Store received joint positions
        self.joint_names = joint_names

        osc_startup()
        osc_udp_server("0.0.0.0", 6080, "osc_server")
        print("Server started on 0.0.0.0:8000 \n ready to receive messages in the following format: /joint_trajectroy [tcp_coordinates0, tcp_coordinates1, ...] optional: timestamp as last element of each tcp_coordinates")
        # Register OSC handler
        osc_method("/joint_trajectory", self.joint_angles_handler, argscheme=osm.OSCARG_DATAUNPACK)

    def joint_angles_handler(self, *args):
        """Handles incoming OSC messages for joint positions."""
        time1 = time.time()
        print("Received joint positions")
        msg = JointTrajectory()
        msg.joint_names = self.joint_names
        joint_positions = [0.0] * len(self.joint_names)
        steps = 30
        vel = 0.4
        if len(args[0]) == len(self.joint_names):
            n=2.0
            for i in range(len(args)-1):
                x, y, z, roll, pitch, yaw = args[i]
                Tep1 = sm.SE3(x, y, z) * sm.SE3.RPY([roll, pitch, yaw], order='xyz')
                x, y, z, roll, pitch, yaw = args[i+1]
                Tep2 = sm.SE3(x, y, z) * sm.SE3.RPY([roll, pitch, yaw], order='xyz')
                cart_traj = rtb.ctraj(Tep1, Tep2, steps)
                for j in range(steps-1):
                    sol = self.robot.ik_LM(cart_traj[j], q0=joint_positions)
                    dist = np.linalg.norm(cart_traj[j].t - cart_traj[j+1].t)
                    point = JointTrajectoryPoint()
                    point.positions = list(sol[0])
                    joint_positions = list(sol[0])
                    point.time_from_start.sec = int(n)
                    point.time_from_start.nanosec = int((n - int(n)) * 1e9)
                    n+=dist/vel
                    n+=0.1
                    msg.points.append(point)

        elif len(args[0]) == len(self.joint_names) + 1:
            for i in range(len(args)):
                x, y, z, roll, pitch, yaw, timetag = args[i]
                Tep = sm.SE3(x, y, z) * sm.SE3.RPY([roll, pitch, yaw], order='xyz')
                x, y, z, roll, pitch, yaw = args[i+1][:-1]
                Tep2 = sm.SE3(x, y, z) * sm.SE3.RPY([roll, pitch, yaw], order='xyz')
                cart_traj = rtb.ctraj(Tep, Tep2, steps)
                for Tep in cart_traj:
                    sol = self.robot.ik_LM(Tep, q0=joint_positions)

        else:
            print("Invalid number or format of arguments")
        
        self.publisher.publish(msg)
        print("published joint positions")
        print(f'Frequency: {round(1/(time.time()-time1),2)} Hz')


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(joint_names, robot)

    # 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()
