import rclpy
from rclpy.node import Node
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from sensor_msgs.msg import JointState
from osc4py3.as_allthreads 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
import os

class ScaledJointTrajectoryPublisher(Node):
    """Node to publish joint trajectories based on OSC messages."""
    def __init__(self, joint_names, robot, cost_mask, joint_velocity_limits):
        super().__init__('scaled_joint_trajectory_publisher')
        
        self.joint_velocity_limits = [joint_velocity_limits[joint] for joint in joint_names]
        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
        )
        self.subscription = self.create_subscription(
            JointState,
            '/joint_states',
            self.joint_states_callback,
            1  # Increased queue size for joint states
            )




        self.maximum_acceleration = [0.0] * len(joint_names)
        # Store received joint positions
        self.joint_names = joint_names
        for joint in joint_names:
            self.maximum_acceleration[joint_names.index(joint)] = float(input(f"Enter the maximum acceleration for joint {joint}: "))
        
        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)
        print("OSC method registered for /joint_trajectory")
    
    def joint_states_callback(self, msg):
        """Callback function to handle incoming joint states."""
        joint_position_dict = dict(zip(msg.name, msg.position))
        self.current_joint_positions = [joint_position_dict[name] for name in self.joint_names]
        joint_velocity_dict = dict(zip(msg.name, msg.velocity))
        self.current_joint_velocities = [joint_velocity_dict[name] for name in self.joint_names]
    
    def joint_angles_handler(self, *args):
        """Handles incoming OSC messages for joint positions."""
        try:
            print("Received joint positions")
            msg = JointTrajectory()
            msg.joint_names = self.joint_names
            steps_per_m = 50

            prev_duration = 0
            for i in range(len(args)-1):
                x, y, z, roll, pitch, yaw = args[i]
                x1, y1, z1, roll1, pitch1, yaw1 = args[i+1]
                steps = int(np.linalg.norm([x1-x, y1-y, z1-z])*steps_per_m)
                if steps <= 1: steps = 2
                cart_traj = [sm.SE3([x+(x1-x)/(steps-1)*i, y+(y1-y)/(steps-1)*i, z+(z1-z)/(steps-1)*i]) * sm.SE3.RPY([roll+(roll1-roll)/(steps-1)*i, pitch+(pitch1-pitch)/(steps-1)*i, yaw+(yaw1-yaw)/(steps-1)*i], order='xyz') for i in range(steps)]
                prev_sol = self.current_joint_positions if i == 0  else sol[0]
                sol_set = []
                for j in (range(steps) if i == 0 else range(1,steps)):
                    sol = self.robot.ik_LM(cart_traj[j], q0=[0.0] * len(self.joint_names), mask = self.cost_mask, joint_limits = True) if i == 0 else self.robot.ik_LM(cart_traj[j], q0=prev_sol, mask = self.cost_mask, joint_limits = True)
                    if sol[1] == 1:
                        sol_set.append(sol[0])
                        prev_sol = list(sol[0])
                    else: print(f'IK could not find a solution for (x,y,z) = ({cart_traj[j].t}), (roll,pitch,yaw) = ({cart_traj[j].rpy()})!')
                distance = abs(sol_set[0]-sol_set[-1])
                ts= distance/np.array(self.joint_velocity_limits)+2*np.array(self.joint_velocity_limits)/np.array(self.maximum_acceleration)
                t = max(ts)
                idx = list(ts).index(t)
                s_acc = self.joint_velocity_limits[idx]**2/(2*self.maximum_acceleration[idx])
                print(f"t: {t}, idx: {idx}, s_acc: {s_acc}")
                print(f"sol_set: {sol_set}")

                for sol in sol_set:
                    print(f"sol: {sol}")
                    s = abs(sol[idx]-sol_set[0][idx])
                    print(f"sol_set[0][idx]: {sol_set[0][idx]}, sol[idx]: {sol[idx]}, s: {s}")
                    if s <= s_acc:
                        duration = np.sqrt(s/self.maximum_acceleration[idx])
                        print(f"acceleration phase, duration: {duration}")
                    elif s <= sol_set[-1][idx]-s_acc:
                        duration = self.joint_velocity_limits[idx]/self.maximum_acceleration[idx] + (s-s_acc)/self.joint_velocity_limits[idx]
                        print(f"constant velocity phase, duration: {duration}")
                    else:
                        duration = t-np.sqrt((sol_set[-1][idx]-s)/self.maximum_acceleration[idx])
                        print(f"deceleration phase, duration: {duration}")
                    
                    point = JointTrajectoryPoint()
                    point.positions = list(sol)
                    duration += prev_duration
                    point.time_from_start.sec = int(duration)
                    point.time_from_start.nanosec = int((duration - int(duration)) * 1e9)
                    msg.points.append(point)
                prev_duration = duration
            msg.header.stamp = self.get_clock().now().to_msg()
            self.publisher.publish(msg)
        except Exception as e:
            print(f"Error in joint angles handler: {e}")

def main():
    """Main function to get joint names and start the ROS 2 & OSC system."""
    while True:
        path_to_urdf = input("Enter the path to the URDF file: ")
        if os.path.isfile(path_to_urdf):
            if not path_to_urdf.endswith('.urdf'):
                print("The file is not a URDF file. Please enter a valid URDF file.")
                continue
            break
        else:
            print("Invalid path. Please enter a valid 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)

    joint_velocity_limits = {}

    # Iterate over all joints in the URDF
    for joint in root.findall('.//joint'):
        joint_name = joint.get('name')  # Get the name of the joint
        
        # Look for the <limit> tag under each joint
        limit = joint.find('limit')
        
        if limit is not None:
            # Extract the velocity limit (if it exists)
            velocity_limit = limit.get('velocity')
            
            if velocity_limit is not None:
                joint_velocity_limits[joint_name] = float(velocity_limit)



    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), of which <= {robot.n} are 1): ")]
            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, joint_velocity_limits)

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