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, cost_mask):
        super().__init__('scaled_joint_trajectory_publisher')
        
        self.cost_mask = cost_mask
        self.robot = robot
        self.trajectroy_topic_name = input("Enter the topic name to which the joint trajectory should be sent to: ")

        if self.trajectroy_topic_name == "":
            self.trajectroy_topic_name = '/scaled_joint_trajectory_controller/joint_trajectory'

        # ROS2 Publisher
        self.publisher = self.create_publisher(
            JointTrajectory, 
            self.trajectroy_topic_name, 
            10
        )

        # Store received joint positions
        self.joint_names = joint_names
        
        self.port = 8000 # UDP port

        osc_startup()
        osc_udp_server("0.0.0.0", self.port, "osc_server")
        print(f"Server started on 0.0.0.0:{str(self.port)} \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 True: #len(args[0]) == len(self.joint_names):
            n=2.0
            for i in range(len(args)-1):
                print(f'i = {i}')
                x, y, z, roll, pitch, yaw = args[i]
                print(1)
                Tep1 = sm.SE3(x, y, z) * sm.SE3.RPY([roll, pitch, yaw], order='xyz')
                print(2)
                x, y, z, roll, pitch, yaw = args[i+1]
                print(3)
                Tep2 = sm.SE3(x, y, z) * sm.SE3.RPY([roll, pitch, yaw], order='xyz')
                print(4)
                cart_traj = rtb.ctraj(Tep1, Tep2, steps)
                print(cart_traj)
                print(5)
                for j in range(steps):
                    print(f'j = {j}')
                    print(6)
                    sol = self.robot.ik_LM(cart_traj[j], q0=joint_positions, mask = self.cost_mask, joint_limits = True)
                    print(7)
                    if sol[1] == 1:
                        print(8)
                        if j == 0: dist = vel*n
                        else: dist = np.linalg.norm(cart_traj[j].t - cart_traj[j-1].t)
                        print(9)
                        point = JointTrajectoryPoint()
                        print(10)
                        point.positions = list(sol[0])
                        print(11)
                        joint_positions = list(sol[0])
                        print(12)
                        point.time_from_start.sec = int(n)
                        print(13)
                        point.time_from_start.nanosec = int((n - int(n)) * 1e9)
                        print(14)
                        n+=dist/vel
                        print(16)
                        msg.points.append(point)
                        print(17)
                    else: print('IK could not find a solution!')
        print(18)
        self.publisher.publish(msg)
        print(19)
        print(f"published joint positions {msg.points[-1]}")
        print(f'Frequency: {round(1/(time.time()-time1),2)} Hz')

'''
        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")'''
    


def main():
    """Main function to get joint names and start the ROS 2 & OSC system."""
    path_to_urdf = input("Enter the path to the URDF file: ")
    tree = ET.parse(path_to_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(path_to_urdf)
    print(robot)
    rclpy.init()
    while True:
        try:
            print("The cost mask determines which coordinates are used for the IK. Each element of the cost mask corresponds to a catesian coordinate [x, y, z, Rx, Ry, Rz].")
            print("The cost mask [1, 1, 1, 0, 0, 0] means that the IK will only consider translation and no rotaion.")
            cost_mask = [int(i) for i in input(f"Enter the cost mask (6 integers (1 or 0) separated by spaces, of which <= {robot.n} are 1): ").split()]
            if sum(cost_mask) <= robot.n and len(cost_mask) == 6:
                break
            else:
                print(f"Invalid input. Expected 6 integers of which {robot.n if robot.n < 6 else 6} or less are 1.")
        except ValueError:
            print("Invalid input. Please enter integers only.")
    print(f"Cost mask: {cost_mask}")
    node = ScaledJointTrajectoryPublisher(joint_names, robot, cost_mask)

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