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
        )
        # 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)
        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")
            viapoints = []
            msg = JointTrajectory()
            msg.joint_names = self.joint_names
            steps_per_m = 4
            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(np.array([x1-x, y1-y, z1-z])) * steps_per_m)
                if steps < 2: 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)]
                if i == 0: prev_sol = self.current_joint_positions
                for j in (range(steps) if i == 0 else range(1,steps)):
                    #print(cart_traj[j])
                    sol = self.robot.ik_LM(cart_traj[j], q0=prev_sol, mask = self.cost_mask, joint_limits = True)
                    if sol[1] == 1:
                        viapoints.append(list(sol[0]))
                        prev_sol = list(sol[0])
                    else: print('IK could not find a solution!')
            dt = 0.01
            tacc = 0.5
            print(f'length viapoints: {len(viapoints)}')
            traj = rtb.mstraj(np.array(viapoints), q0 = self.current_joint_positions ,dt=dt, tacc=tacc, qdmax=[1 * i for i in self.joint_velocity_limits])
            print(len(traj.q))
            print(len(traj.t))
            print(traj.t)
            print(traj.arrive)
            msg.points = []
            for i in range(len(traj.q)):
                point = JointTrajectoryPoint()
                point.positions = list(traj.q[i])
                point.time_from_start.sec = int(traj.t[i])
                point.time_from_start.nanosec = int(((traj.t[i] - int(traj.t[i])) * 1e9))
                #point.time_from_start = rclpy.duration.Duration(seconds=traj.t[i]).to_msg()
                msg.points.append(point)
            msg.header.stamp = self.get_clock().now().to_msg()
            self.publisher.publish(msg)
            print('published')
        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()
