import xml.etree.ElementTree as ET
import roboticstoolbox as rtb

def main():
    """Main function to get joint names and start the ROS 2 & OSC system."""
    robot_urdf = input("Enter the path to the URDF file: ")
    tree = ET.parse(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(robot_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)
    
    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}")

    print(f'robot: {robot}')
    print(f'joint names: {joint_names}')
    print(f'joint velocity limits: {joint_velocity_limits}')


if __name__ == '__main__':
    main()