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 time
import numpy as np
import spatialmath as sm
import roboticstoolbox as rtb
import os

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

        self.subscription = self.create_subscription(
            JointState,
            '/joint_states',
            self.joint_states_callback,
            1  # Increased queue size for joint states
        )

        # Store received joint positions
        self.current_joint_positions = [0.0] * len(joint_names)
        self.joint_names = joint_names
        self.joint_velocity_limits = joint_velocity_limits
        self.desired_joint_positions = [0.0] * len(joint_names)
        self.previous_desired = [0.0] * len(joint_names)
        self.robot = robot
        self.cost_mask = cost_mask
        
        ip = "0.0.0.0"  # Listen on all network interfaces
        port = 8000     # Must match the sender's port in `joint_state_osc.py`

        osc_startup()
        osc_udp_server(ip, port, "osc_server")
        print("Server started on 0.0.0.0:8000 \n ready to receive messages in the following format: /joint_trajectroy [joint_positions] optional: duration as last argument")
        # Register OSC handler
        osc_method("/joint_angles", self.joint_angles_handler, argscheme=osm.OSCARG_DATAUNPACK)


        self.hz = float(input("Enter the desired refresh frequency (Hz): "))
        # Start the OSC server in a separate thread to avoid blocking the ROS 2 event loop
        self.create_timer(1/self.hz, self.update_position)  # Timer to call osc_process periodically

        while True:
            try:
                self.x_limits = [float(i) if i != 'x' else None for i in input("Enter the lower and upper limits for x (space-separated, enter 'x' for no limit): ").split()]
                self.y_limits = [float(i) if i != 'x' else None for i in input("Enter the lower and upper limits for y (space-separated, enter 'x' for no limit): ").split()]
                self.z_limits = [float(i) if i != 'x' else None for i in input("Enter the lower and upper limits for z (space-separated, enter 'x' for no limit): ").split()]
                
                if len(self.x_limits) != 2 or len(self.y_limits) != 2 or len(self.z_limits) != 2:
                    print("Invalid input. Please enter exactly two values (or leave blank) for each limit.")
                    continue
                
                if (self.x_limits[0] is not None and self.x_limits[1] is not None and self.x_limits[0] >= self.x_limits[1]) or \
                (self.y_limits[0] is not None and self.y_limits[1] is not None and self.y_limits[0] >= self.y_limits[1]) or \
                (self.z_limits[0] is not None and self.z_limits[1] is not None and self.z_limits[0] >= self.z_limits[1]):
                    print("Invalid input. Lower limit must be less than upper limit for each axis.")
                    continue
                
                print(f"Current limits:")
                print(f"x: {self.x_limits}")
                print(f"y: {self.y_limits}")
                print(f"z: {self.z_limits}")
                confirm = input("Do you want your robot to move in this range? (y/n): ").strip().lower()
                if confirm == 'y':
                    break
                elif confirm == 'n':
                    print("Please re-enter the limits.")
                else:
                    print("Invalid input. Please enter 'y' or 'n'.")
            except ValueError:
                print("Invalid input. Please enter numeric values only.")

    def joint_angles_handler(self, *args):
    # Ensure the desired joint positions are within the specified limits

        x, y, z, r, p, yaw, *_ = [float(i) for i in list(args)]
        if self.x_limits[0] is not None:
            x = max(self.x_limits[0], x)
        if self.x_limits[1] is not None:
            x = min(self.x_limits[1], x)
        if self.y_limits[0] is not None:
            y = max(self.y_limits[0], y)
        if self.y_limits[1] is not None:
            y = min(self.y_limits[1], y)
        if self.z_limits[0] is not None:
            z = max(self.z_limits[0], z)
        if self.z_limits[1] is not None:
            z = min(self.z_limits[1], z)

        if x != args[0] or y != args[1] or z != args[2]:
            self.get_logger().warn(
            f"Desired joint positions adjusted to fit within limits: "
            f"x={x}, y={y}, z={z} (original: x={args[0]}, y={args[1]}, z={args[2]})"
            )

        self.desired_joint_positions = [x, y, z, r, p, yaw]
        

    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]

    def rampfunction(self, startvalue, blendtime, currenttime):
        """
        Ramp function to create a smooth transition from startvalue to 1 over blendtime seconds.
        """
        if currenttime < blendtime:
            return startvalue + (1 - startvalue) * (currenttime / blendtime)
        else:
            return 1

    def update_position(self):
        if self.desired_joint_positions == self.previous_desired:
            return
        msg = JointTrajectory()
        msg.joint_names = self.joint_names
        steps_per_m = 30
        if True: #len(args[0]) == len(self.joint_names):
            prev_duration = 0
            T1 = self.robot.fkine(self.current_joint_positions)
            [x,y,z] = T1.t
            [roll, pitch, yaw] = T1.rpy()
            x1, y1, z1, roll1, pitch1, yaw1 = self.desired_joint_positions
            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]) for i in range(steps)]
            for j in range(steps):
                sol = self.robot.ik_LM(cart_traj[j], q0=self.current_joint_positions, mask = self.cost_mask, joint_limits = True)
                if sol[1] == 1:
                    duration = 0
                    prev = self.current_joint_positions if j == 0 else prev_sol
                    for p1, p2, max_vel in zip(sol[0], prev, self.joint_velocity_limits.values()):
                        duration = max(duration, abs(p1 - p2) / max_vel)#, 1/self.hz) # as minimun
                    point = JointTrajectoryPoint()
                    point.positions = list(sol[0])
                    duration /= self.rampfunction(0.1, 2, prev_duration)
                    duration += prev_duration
                    prev_duration = duration
                    point.time_from_start.sec = int(duration)
                    point.time_from_start.nanosec = int((duration - int(duration)) * 1e9)
                    msg.points.append(point)
                    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()})!')
                    prev_sol = self.current_joint_positions
        msg.header.stamp = self.get_clock().now().to_msg()
        self.publisher.publish(msg)
        self.previous_desired = self.desired_joint_positions

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)
    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)
    
    print("The cost mask determines which coordinates are used for the IK. Each element of the cost mask corresponds to a Cartesian coordinate [x, y, z, roll, pitch, yaw].")
    print("The cost mask [1, 1, 1, 0, 0, 0] means that the IK will only consider translation and no rotation.")
    while True:
        try:
            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:
                considered_coords = [coord for coord, use in zip(['x', 'y', 'z', 'roll', 'pitch', 'yaw'], cost_mask) if use == 1]
                print(f"The following coordinates will be considered for IK: {', '.join(considered_coords)}")
                confirm = input("Are you sure you want to proceed with this cost mask? (y/n): ").strip().lower()
                if confirm == 'y':
                    break
                elif confirm == 'n':
                    print("Please re-enter the cost mask.")
                else:
                    print("Invalid input. Please enter 'y' or 'n'.")
            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}")

    rclpy.init()

    node = ScaledJointTrajectoryPublisher(joint_names, joint_velocity_limits, robot, cost_mask)
    
    # Run ROS 2 spin, and osc_process will be handled by the timer
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        print("")
    finally:
        node.destroy_node()
        rclpy.shutdown()
        osc_terminate()

if __name__ == '__main__':
    main()
