RoboticArmTest/src/robot/kinematics.py

#!/usr/bin/env python3
"""
KDL-based Kinematics Engine

Provides forward and inverse kinematics calculations using KDL library.
Supports URDF-based robot configurations for precise kinematic computations.
"""

import PyKDL as kdl
import numpy as np
from typing import List, Tuple, Optional, Dict
from urdf_parser_py.urdf import URDF
import sys
import os


class KinematicsEngine:
    """KDL-based kinematics engine for robot arm calculations"""
    
    def __init__(self, urdf_path: str, base_link: str, end_link: str, 
                 base_position: List[float] = None, base_orientation: List[float] = None):
        """Initialize KDL kinematics engine from URDF file"""
        self.urdf_path = urdf_path
        self.base_link = base_link
        self.end_link = end_link
        
        # Base coordinate transformation (base pose in world frame)
        self.base_position = base_position if base_position else [0.0, 0.0, 0.0]
        self.base_orientation = base_orientation if base_orientation else [0.0, 0.0, 0.0]  # RPY
        self._base_to_world_frame = self._create_transform_frame(self.base_position, self.base_orientation)
        self._world_to_base_frame = self._base_to_world_frame.Inverse()
        
        # KDL components
        self.chain: Optional[kdl.Chain] = None
        self.fk_solver: Optional[kdl.ChainFkSolverPos_recursive] = None
        self.ik_solver: Optional[kdl.ChainIkSolverPos_NR_JL] = None
        self.jac_solver: Optional[kdl.ChainJntToJacSolver] = None
        self.joint_limits: List[Tuple[float, float]] = []
        
        # Initialize KDL chain and solvers
        self._build_kdl_chain()
        self._create_solvers()
    
    def _build_kdl_chain(self) -> None:
        """Build KDL chain from URDF file"""
        try:
            # Parse URDF
            robot = URDF.from_xml_file(self.urdf_path)
            
            # Build chain from base to end effector
            self.chain = self._urdf_to_kdl_chain(robot, self.base_link, self.end_link)
            
            if self.chain is None:
                raise RuntimeError(f"Failed to build KDL chain from {self.base_link} to {self.end_link}")
                
        except Exception as e:
            raise RuntimeError(f"Failed to build KDL chain: {e}")
    
    def _urdf_to_kdl_chain(self, robot: URDF, base_link: str, end_link: str) -> kdl.Chain:
        """Convert URDF robot to KDL chain"""
        chain = kdl.Chain()
        
        # Find path from base to end link
        path = self._find_link_path(robot, base_link, end_link)
        if not path:
            raise RuntimeError(f"No path found from {base_link} to {end_link}")
        
        # Build chain along the path
        for i in range(len(path) - 1):
            current_link = path[i]
            next_link = path[i + 1]
            
            # Find joint connecting these links
            joint = self._find_joint_between_links(robot, current_link, next_link)
            if joint is None:
                continue
            
            # Convert joint to KDL segment
            segment = self._joint_to_kdl_segment(joint, robot)
            chain.addSegment(segment)
            
            # Store joint limits for controllable joints
            if joint.type in ['revolute', 'prismatic']:
                if joint.limit:
                    self.joint_limits.append((joint.limit.lower, joint.limit.upper))
                else:
                    # Default large limits if not specified
                    self.joint_limits.append((-3.14159, 3.14159))
        
        return chain
    
    def _find_link_path(self, robot: URDF, start_link: str, end_link: str) -> List[str]:
        """Find path between two links in URDF tree"""
        if start_link == end_link:
            return [start_link]
        
        # Build parent-child relationships
        parent_map = {}
        for joint in robot.joints:
            parent_map[joint.child] = joint.parent
        
        # Find path from end to start
        path = [end_link]
        current = end_link
        
        while current != start_link and current in parent_map:
            current = parent_map[current]
            path.append(current)
        
        if current != start_link:
            return []  # No path found
        
        # Reverse to get start->end path
        return list(reversed(path))
    
    def _find_joint_between_links(self, robot: URDF, parent_link: str, child_link: str):
        """Find joint connecting two links"""
        for joint in robot.joints:
            if joint.parent == parent_link and joint.child == child_link:
                return joint
        return None
    
    def _joint_to_kdl_segment(self, joint, robot: URDF) -> kdl.Segment:
        """Convert URDF joint to KDL segment"""
        # Get joint origin
        origin = joint.origin if joint.origin else URDF.Origin()
        xyz = origin.xyz if origin.xyz else [0, 0, 0]
        rpy = origin.rpy if origin.rpy else [0, 0, 0]
        
        # Create KDL frame
        rotation = kdl.Rotation.RPY(rpy[0], rpy[1], rpy[2])
        translation = kdl.Vector(xyz[0], xyz[1], xyz[2])
        frame = kdl.Frame(rotation, translation)
        
        # Create KDL joint
        if joint.type == 'revolute':
            axis = joint.axis if joint.axis else [0, 0, 1]
            kdl_joint = kdl.Joint(joint.name, kdl.Joint.RotZ if axis == [0, 0, 1] 
                                 else kdl.Joint.RotY if axis == [0, 1, 0] 
                                 else kdl.Joint.RotX)
        elif joint.type == 'prismatic':
            axis = joint.axis if joint.axis else [0, 0, 1]
            kdl_joint = kdl.Joint(joint.name, kdl.Joint.TransZ if axis == [0, 0, 1]
                                 else kdl.Joint.TransY if axis == [0, 1, 0]
                                 else kdl.Joint.TransX)
        else:
            # Fixed joint
            kdl_joint = kdl.Joint(joint.name, kdl.Joint.None)
        
        return kdl.Segment(joint.child, kdl_joint, frame)
    
    def _create_solvers(self) -> None:
        """Create KDL solvers for forward kinematics, inverse kinematics and Jacobian"""
        if self.chain is None:
            raise RuntimeError("KDL chain not initialized")
        
        try:
            # Forward kinematics solver
            self.fk_solver = kdl.ChainFkSolverPos_recursive(self.chain)
            
            # Jacobian solver
            self.jac_solver = kdl.ChainJntToJacSolver(self.chain)
            
            # Inverse kinematics solver with joint limits
            if self.joint_limits:
                q_min = kdl.JntArray(len(self.joint_limits))
                q_max = kdl.JntArray(len(self.joint_limits))
                
                for i, (lower, upper) in enumerate(self.joint_limits):
                    q_min[i] = lower
                    q_max[i] = upper
                
                # Create velocity IK solver first
                ik_vel_solver = kdl.ChainIkSolverVel_pinv(self.chain)
                
                # Create position IK solver with joint limits
                self.ik_solver = kdl.ChainIkSolverPos_NR_JL(
                    self.chain, q_min, q_max, self.fk_solver, ik_vel_solver
                )
            else:
                # Fallback without joint limits
                ik_vel_solver = kdl.ChainIkSolverVel_pinv(self.chain)
                self.ik_solver = kdl.ChainIkSolverPos_NR(
                    self.chain, self.fk_solver, ik_vel_solver
                )
                
        except Exception as e:
            raise RuntimeError(f"Failed to create KDL solvers: {e}")
    
    def forward_kinematics(self, joint_positions: List[float]) -> Tuple[List[float], List[float]]:
        """Calculate forward kinematics from joint positions to world coordinate end effector pose"""
        if self.fk_solver is None:
            raise RuntimeError("Forward kinematics solver not initialized")
        
        if len(joint_positions) != self.chain.getNrOfJoints():
            raise ValueError(f"Expected {self.chain.getNrOfJoints()} joint positions, got {len(joint_positions)}")
        
        # Create joint array
        q = kdl.JntArray(len(joint_positions))
        for i, pos in enumerate(joint_positions):
            q[i] = pos
        
        # Solve forward kinematics in base frame
        base_end_frame = kdl.Frame()
        result = self.fk_solver.JntToCart(q, base_end_frame)
        
        if result != 0:
            raise RuntimeError("Forward kinematics calculation failed")
        
        # Transform from base frame to world frame
        world_end_frame = self._base_to_world_frame * base_end_frame
        
        # Extract world position
        position = [
            world_end_frame.p.x(),
            world_end_frame.p.y(), 
            world_end_frame.p.z()
        ]
        
        # Extract world orientation as quaternion
        rotation = world_end_frame.M
        orientation = self._rotation_to_quaternion(rotation)
        
        return position, orientation
    
    def inverse_kinematics(self, target_position: List[float], 
                          target_orientation: List[float] = None,
                          seed_angles: List[float] = None) -> Optional[List[float]]:
        """Calculate inverse kinematics from world coordinate target pose to joint positions"""
        if self.ik_solver is None:
            raise RuntimeError("Inverse kinematics solver not initialized")
        
        # Create world target frame
        world_target_frame = kdl.Frame()
        world_target_frame.p = kdl.Vector(target_position[0], target_position[1], target_position[2])
        
        if target_orientation is not None:
            world_target_frame.M = self._quaternion_to_rotation(target_orientation)
        
        # Transform from world frame to base frame
        base_target_frame = self._world_to_base_frame * world_target_frame
        
        # Set seed angles
        q_init = kdl.JntArray(self.chain.getNrOfJoints())
        if seed_angles:
            if len(seed_angles) != self.chain.getNrOfJoints():
                raise ValueError(f"Expected {self.chain.getNrOfJoints()} seed angles, got {len(seed_angles)}")
            for i, angle in enumerate(seed_angles):
                q_init[i] = angle
        else:
            # Use zero as default seed
            for i in range(self.chain.getNrOfJoints()):
                q_init[i] = 0.0
        
        # Solve inverse kinematics in base frame
        q_out = kdl.JntArray(self.chain.getNrOfJoints())
        result = self.ik_solver.CartToJnt(q_init, base_target_frame, q_out)
        
        if result < 0:
            return None  # No solution found
        
        # Return joint angles
        joint_angles = []
        for i in range(q_out.rows()):
            joint_angles.append(q_out[i])
        
        return joint_angles
    
    def compute_jacobian(self, joint_positions: List[float]) -> np.ndarray:
        """Compute Jacobian matrix for given joint configuration"""
        if self.jac_solver is None:
            raise RuntimeError("Jacobian solver not initialized")
        
        if len(joint_positions) != self.chain.getNrOfJoints():
            raise ValueError(f"Expected {self.chain.getNrOfJoints()} joint positions, got {len(joint_positions)}")
        
        # Create joint array
        q = kdl.JntArray(len(joint_positions))
        for i, pos in enumerate(joint_positions):
            q[i] = pos
        
        # Compute Jacobian
        jacobian = kdl.Jacobian(self.chain.getNrOfJoints())
        result = self.jac_solver.JntToJac(q, jacobian)
        
        if result != 0:
            raise RuntimeError("Jacobian calculation failed")
        
        # Convert to numpy array
        jac_array = np.zeros((6, self.chain.getNrOfJoints()))
        for i in range(6):
            for j in range(self.chain.getNrOfJoints()):
                jac_array[i, j] = jacobian[i, j]
        
        return jac_array
    
    def get_joint_limits(self) -> List[Tuple[float, float]]:
        """Get joint limits for all joints in the chain"""
        return self.joint_limits.copy()
    
    def validate_joint_configuration(self, joint_positions: List[float]) -> bool:
        """Validate if joint configuration is within limits"""
        if len(joint_positions) != len(self.joint_limits):
            return False
        
        for i, (pos, (lower, upper)) in enumerate(zip(joint_positions, self.joint_limits)):
            if pos < lower or pos > upper:
                return False
        
        return True
    
    def get_num_joints(self) -> int:
        """Get number of joints in the kinematic chain"""
        return self.chain.getNrOfJoints() if self.chain else 0
    
    def _rotation_to_quaternion(self, rotation: kdl.Rotation) -> List[float]:
        """Convert KDL rotation to quaternion [x, y, z, w]"""
        # Get quaternion from KDL rotation
        quat = rotation.GetQuaternion()
        return [quat[0], quat[1], quat[2], quat[3]]  # [x, y, z, w]
    
    def _quaternion_to_rotation(self, quaternion: List[float]) -> kdl.Rotation:
        """Convert quaternion [x, y, z, w] to KDL rotation"""
        return kdl.Rotation.Quaternion(quaternion[0], quaternion[1], quaternion[2], quaternion[3])
    
    def _create_transform_frame(self, position: List[float], orientation_rpy: List[float]) -> kdl.Frame:
        """Create KDL frame from position and RPY orientation"""
        rotation = kdl.Rotation.RPY(orientation_rpy[0], orientation_rpy[1], orientation_rpy[2])
        translation = kdl.Vector(position[0], position[1], position[2])
        return kdl.Frame(rotation, translation)


def create_kinematics_engine(config_loader) -> KinematicsEngine:
    """Create KinematicsEngine from configuration loader"""
    robot_config = config_loader.get_robot_config()
    urdf_path = robot_config['model_path']
    
    # Get robot base pose from config
    base_position = robot_config.get('base_position', [0.0, 0.0, 0.0])
    base_orientation = robot_config.get('base_orientation', [0.0, 0.0, 0.0])
    
    # Get link names from config or use defaults
    kinematics_config = robot_config.get('kinematics', {})
    base_link = kinematics_config.get('base_link', 'base_link')
    end_link = kinematics_config.get('end_link', 'link_9')
    
    return KinematicsEngine(urdf_path, base_link, end_link, base_position, base_orientation)