unity2moveit2/unity-project/Assets/Scripts/Core/UniversalRobotAdapter.cs

using System;
using System.Collections.Generic;
using UnityEngine;

namespace UnityMoveIt2.Core
{
    /// <summary>
    /// 通用机器人适配器 - Universal Robot Adapter
    /// 支持1-9自由度的任意机械臂配置 / Supports 1-9 DOF arbitrary manipulator configuration
    /// </summary>
    [RequireComponent(typeof(ArticulationBody))]
    public class UniversalRobotAdapter : MonoBehaviour
    {
        #region Inspector Fields

        [Header("机器人配置 / Robot Configuration")]
        [SerializeField, Range(1, 9)]
        [Tooltip("机器人自由度数量 / Degrees of Freedom")]
        private int degreesOfFreedom = 6;

        [SerializeField]
        [Tooltip("机器人名称 / Robot Name")]
        private string robotName = "universal_robot";

        [SerializeField]
        [Tooltip("机器人类型 / Robot Type")]
        private string robotType = "serial_manipulator";

        [Header("坐标系配置 / Frame Configuration")]
        [SerializeField]
        [Tooltip("基座坐标系 / Base Link")]
        private Transform baseLink;

        [SerializeField]
        [Tooltip("末端执行器 / End Effector")]
        private Transform endEffector;

        [SerializeField]
        [Tooltip("工具中心点偏移 / TCP Offset")]
        private Vector3 tcpOffset = Vector3.zero;

        [Header("URDF配置 / URDF Configuration")]
        [SerializeField]
        [Tooltip("URDF文件路径 / URDF Path")]
        private string urdfPath;

        [SerializeField]
        [Tooltip("SRDF文件路径 / SRDF Path")]
        private string srdfPath;

        [Header("运动学参数 / Kinematic Parameters")]
        [SerializeField]
        [Tooltip("工作空间半径(米) / Workspace Radius (m)")]
        private float workspaceRadius = 1.0f;

        [SerializeField]
        [Tooltip("规划组名称 / Planning Group Name")]
        private string planningGroup = "manipulator";

        #endregion

        #region Private Fields

        private ArticulationBody[] joints;
        private JointData[] jointData;
        private bool isInitialized = false;
        private Transform[] linkTransforms;
        private Dictionary<string, int> jointNameToIndex;

        #endregion

        #region Properties

        /// <summary>
        /// 自由度数量 / Degrees of Freedom
        /// </summary>
        public int DOF => degreesOfFreedom;

        /// <summary>
        /// 机器人名称 / Robot Name
        /// </summary>
        public string RobotName => robotName;

        /// <summary>
        /// 是否已初始化 / Is Initialized
        /// </summary>
        public bool IsInitialized => isInitialized;

        /// <summary>
        /// 基座变换 / Base Transform
        /// </summary>
        public Transform BaseLink => baseLink;

        /// <summary>
        /// 末端执行器变换 / End Effector Transform
        /// </summary>
        public Transform EndEffector => endEffector;

        /// <summary>
        /// 关节数据数组 / Joint Data Array
        /// </summary>
        public JointData[] JointData => jointData;

        /// <summary>
        /// 规划组名称 / Planning Group Name
        /// </summary>
        public string PlanningGroup => planningGroup;

        #endregion

        #region Unity Lifecycle

        private void Awake()
        {
            InitializeRobot();
        }

        private void Start()
        {
            if (!isInitialized)
            {
                Debug.LogWarning($"[UniversalRobotAdapter] {robotName} 未能正确初始化 / Failed to initialize properly");
            }
        }

        #endregion

        #region Initialization

        /// <summary>
        /// 初始化机器人 / Initialize Robot
        /// </summary>
        public bool InitializeRobot()
        {
            try
            {
                Debug.Log($"[UniversalRobotAdapter] 初始化 {degreesOfFreedom} DOF 机器人: {robotName}");

                // 验证基本配置 / Validate basic configuration
                if (!ValidateConfiguration())
                {
                    Debug.LogError("[UniversalRobotAdapter] 配置验证失败 / Configuration validation failed");
                    return false;
                }

                // 查找所有关节 / Find all joints
                DiscoverJoints();

                // 初始化关节数据 / Initialize joint data
                InitializeJointData();

                // 构建关节索引映射 / Build joint name to index mapping
                BuildJointMapping();

                // 设置默认关节位置 / Set default joint positions
                SetDefaultJointPositions();

                isInitialized = true;
                Debug.Log($"[UniversalRobotAdapter] {robotName} 初始化成功 / Initialized successfully");

                return true;
            }
            catch (Exception ex)
            {
                Debug.LogError($"[UniversalRobotAdapter] 初始化失败 / Initialization failed: {ex.Message}");
                return false;
            }
        }

        /// <summary>
        /// 验证配置 / Validate Configuration
        /// </summary>
        private bool ValidateConfiguration()
        {
            if (baseLink == null)
            {
                Debug.LogError("[UniversalRobotAdapter] 基座链接未设置 / Base link not set");
                return false;
            }

            if (endEffector == null)
            {
                Debug.LogWarning("[UniversalRobotAdapter] 末端执行器未设置，尝试自动检测 / End effector not set, attempting auto-detection");
                AutoDetectEndEffector();
            }

            if (degreesOfFreedom < 1 || degreesOfFreedom > 9)
            {
                Debug.LogError($"[UniversalRobotAdapter] 自由度超出范围: {degreesOfFreedom} / DOF out of range");
                return false;
            }

            return true;
        }

        /// <summary>
        /// 发现所有关节 / Discover All Joints
        /// </summary>
        private void DiscoverJoints()
        {
            // 获取所有ArticulationBody组件 / Get all ArticulationBody components
            List<ArticulationBody> foundJoints = new List<ArticulationBody>();
            GetComponentsInChildren(true, foundJoints);

            // 过滤出可移动关节 / Filter movable joints
            List<ArticulationBody> movableJoints = new List<ArticulationBody>();
            foreach (var joint in foundJoints)
            {
                if (joint.jointType != ArticulationJointType.FixedJoint && joint != GetComponent<ArticulationBody>())
                {
                    movableJoints.Add(joint);
                }
            }

            // 验证关节数量 / Validate joint count
            if (movableJoints.Count != degreesOfFreedom)
            {
                Debug.LogWarning($"[UniversalRobotAdapter] 发现关节数({movableJoints.Count})与配置DOF({degreesOfFreedom})不匹配 / " +
                    $"Found joints count does not match configured DOF");
                degreesOfFreedom = movableJoints.Count;
            }

            joints = movableJoints.ToArray();
            Debug.Log($"[UniversalRobotAdapter] 发现 {joints.Length} 个可移动关节 / Discovered {joints.Length} movable joints");
        }

        /// <summary>
        /// 初始化关节数据 / Initialize Joint Data
        /// </summary>
        private void InitializeJointData()
        {
            jointData = new JointData[joints.Length];
            linkTransforms = new Transform[joints.Length];

            for (int i = 0; i < joints.Length; i++)
            {
                jointData[i] = new JointData
                {
                    index = i,
                    name = joints[i].name,
                    jointType = joints[i].jointType,
                    lowerLimit = GetJointLowerLimit(joints[i]),
                    upperLimit = GetJointUpperLimit(joints[i]),
                    maxVelocity = GetJointMaxVelocity(joints[i]),
                    maxAcceleration = GetJointMaxAcceleration(joints[i]),
                    currentPosition = 0f,
                    currentVelocity = 0f,
                    targetPosition = 0f
                };

                linkTransforms[i] = joints[i].transform;

                Debug.Log($"[UniversalRobotAdapter] 关节 {i}: {jointData[i].name}, " +
                    $"范围: [{jointData[i].lowerLimit:F2}, {jointData[i].upperLimit:F2}]");
            }
        }

        /// <summary>
        /// 构建关节名称到索引的映射 / Build Joint Name to Index Mapping
        /// </summary>
        private void BuildJointMapping()
        {
            jointNameToIndex = new Dictionary<string, int>();
            for (int i = 0; i < jointData.Length; i++)
            {
                jointNameToIndex[jointData[i].name] = i;
            }
        }

        /// <summary>
        /// 自动检测末端执行器 / Auto Detect End Effector
        /// </summary>
        private void AutoDetectEndEffector()
        {
            if (joints != null && joints.Length > 0)
            {
                // 使用最后一个关节作为末端执行器 / Use last joint as end effector
                ArticulationBody lastJoint = joints[joints.Length - 1];

                // 查找该关节的子物体作为末端执行器 / Find child of last joint as end effector
                if (lastJoint.transform.childCount > 0)
                {
                    endEffector = lastJoint.transform.GetChild(0);
                    Debug.Log($"[UniversalRobotAdapter] 自动检测到末端执行器: {endEffector.name} / Auto-detected end effector");
                }
                else
                {
                    endEffector = lastJoint.transform;
                    Debug.LogWarning("[UniversalRobotAdapter] 使用最后关节作为末端执行器 / Using last joint as end effector");
                }
            }
        }

        #endregion

        #region Joint Control

        /// <summary>
        /// 设置关节位置 / Set Joint Positions
        /// </summary>
        public void SetJointPositions(float[] positions)
        {
            if (!isInitialized)
            {
                Debug.LogWarning("[UniversalRobotAdapter] 机器人未初始化 / Robot not initialized");
                return;
            }

            if (positions == null || positions.Length != degreesOfFreedom)
            {
                Debug.LogError($"[UniversalRobotAdapter] 关节位置数组长度不匹配: {positions?.Length} vs {degreesOfFreedom}");
                return;
            }

            for (int i = 0; i < positions.Length; i++)
            {
                SetJointPosition(i, positions[i]);
            }
        }

        /// <summary>
        /// 设置单个关节位置 / Set Single Joint Position
        /// </summary>
        public void SetJointPosition(int jointIndex, float position)
        {
            if (jointIndex < 0 || jointIndex >= joints.Length)
            {
                Debug.LogError($"[UniversalRobotAdapter] 关节索引越界: {jointIndex}");
                return;
            }

            // 限制关节位置在允许范围内 / Clamp position within limits
            float clampedPosition = Mathf.Clamp(position, jointData[jointIndex].lowerLimit, jointData[jointIndex].upperLimit);

            if (Mathf.Abs(clampedPosition - position) > 0.001f)
            {
                Debug.LogWarning($"[UniversalRobotAdapter] 关节 {jointIndex} 位置被限制: {position} -> {clampedPosition}");
            }

            // 更新关节目标位置 / Update joint target position
            ArticulationDrive drive = joints[jointIndex].xDrive;
            drive.target = clampedPosition * Mathf.Rad2Deg; // Unity uses degrees
            joints[jointIndex].xDrive = drive;

            jointData[jointIndex].targetPosition = clampedPosition;
        }

        /// <summary>
        /// 获取当前关节位置 / Get Current Joint Positions
        /// </summary>
        public float[] GetCurrentJointPositions()
        {
            float[] positions = new float[joints.Length];
            for (int i = 0; i < joints.Length; i++)
            {
                positions[i] = joints[i].jointPosition[0] * Mathf.Deg2Rad; // Convert to radians
                jointData[i].currentPosition = positions[i];
            }
            return positions;
        }

        /// <summary>
        /// 获取当前关节速度 / Get Current Joint Velocities
        /// </summary>
        public float[] GetCurrentJointVelocities()
        {
            float[] velocities = new float[joints.Length];
            for (int i = 0; i < joints.Length; i++)
            {
                velocities[i] = joints[i].jointVelocity[0] * Mathf.Deg2Rad; // Convert to rad/s
                jointData[i].currentVelocity = velocities[i];
            }
            return velocities;
        }

        /// <summary>
        /// 设置默认关节位置 / Set Default Joint Positions
        /// </summary>
        private void SetDefaultJointPositions()
        {
            float[] zeroPositions = new float[degreesOfFreedom];
            SetJointPositions(zeroPositions);
        }

        #endregion

        #region Forward Kinematics

        /// <summary>
        /// 计算末端执行器位姿(正运动学) / Calculate End Effector Pose (Forward Kinematics)
        /// </summary>
        public Pose GetEndEffectorPose()
        {
            if (endEffector == null)
            {
                Debug.LogError("[UniversalRobotAdapter] 末端执行器未设置 / End effector not set");
                return new Pose(Vector3.zero, Quaternion.identity);
            }

            // 应用TCP偏移 / Apply TCP offset
            Vector3 position = endEffector.position + endEffector.TransformDirection(tcpOffset);
            Quaternion rotation = endEffector.rotation;

            return new Pose(position, rotation);
        }

        /// <summary>
        /// 获取指定链接的位姿 / Get Pose of Specified Link
        /// </summary>
        public Pose GetLinkPose(int linkIndex)
        {
            if (linkIndex < 0 || linkIndex >= linkTransforms.Length)
            {
                Debug.LogError($"[UniversalRobotAdapter] 链接索引越界: {linkIndex}");
                return new Pose(Vector3.zero, Quaternion.identity);
            }

            return new Pose(linkTransforms[linkIndex].position, linkTransforms[linkIndex].rotation);
        }

        #endregion

        #region Joint Limits

        private float GetJointLowerLimit(ArticulationBody joint)
        {
            if (joint.jointType == ArticulationJointType.RevoluteJoint)
            {
                return joint.xDrive.lowerLimit * Mathf.Deg2Rad;
            }
            else if (joint.jointType == ArticulationJointType.PrismaticJoint)
            {
                return joint.xDrive.lowerLimit;
            }
            return -Mathf.PI;
        }

        private float GetJointUpperLimit(ArticulationBody joint)
        {
            if (joint.jointType == ArticulationJointType.RevoluteJoint)
            {
                return joint.xDrive.upperLimit * Mathf.Deg2Rad;
            }
            else if (joint.jointType == ArticulationJointType.PrismaticJoint)
            {
                return joint.xDrive.upperLimit;
            }
            return Mathf.PI;
        }

        private float GetJointMaxVelocity(ArticulationBody joint)
        {
            // Unity ArticulationBody doesn't directly expose max velocity
            // Use a reasonable default based on joint type
            return joint.jointType == ArticulationJointType.RevoluteJoint ? 2.0f : 1.0f;
        }

        private float GetJointMaxAcceleration(ArticulationBody joint)
        {
            // Unity ArticulationBody doesn't directly expose max acceleration
            // Use a reasonable default
            return joint.jointType == ArticulationJointType.RevoluteJoint ? 4.0f : 2.0f;
        }

        #endregion

        #region Utility Methods

        /// <summary>
        /// 根据名称获取关节索引 / Get Joint Index by Name
        /// </summary>
        public int GetJointIndex(string jointName)
        {
            if (jointNameToIndex != null && jointNameToIndex.ContainsKey(jointName))
            {
                return jointNameToIndex[jointName];
            }
            return -1;
        }

        /// <summary>
        /// 检查关节位置是否在限制范围内 / Check if Joint Position is Within Limits
        /// </summary>
        public bool IsJointPositionValid(int jointIndex, float position)
        {
            if (jointIndex < 0 || jointIndex >= jointData.Length)
                return false;

            return position >= jointData[jointIndex].lowerLimit && position <= jointData[jointIndex].upperLimit;
        }

        /// <summary>
        /// 检查所有关节位置是否有效 / Check if All Joint Positions are Valid
        /// </summary>
        public bool AreJointPositionsValid(float[] positions)
        {
            if (positions == null || positions.Length != degreesOfFreedom)
                return false;

            for (int i = 0; i < positions.Length; i++)
            {
                if (!IsJointPositionValid(i, positions[i]))
                    return false;
            }
            return true;
        }

        #endregion

        #region Configuration Methods

        /// <summary>
        /// 从配置文件加载机器人 / Load Robot from Configuration
        /// </summary>
        public bool LoadFromConfiguration(RobotConfiguration config)
        {
            if (config == null)
            {
                Debug.LogError("[UniversalRobotAdapter] 配置为空 / Configuration is null");
                return false;
            }

            robotName = config.robotName;
            degreesOfFreedom = config.dof;
            urdfPath = config.urdfPath;
            srdfPath = config.srdfPath;
            planningGroup = config.planningGroup;

            return InitializeRobot();
        }

        /// <summary>
        /// 导出当前配置 / Export Current Configuration
        /// </summary>
        public RobotConfiguration ExportConfiguration()
        {
            return new RobotConfiguration
            {
                robotName = this.robotName,
                dof = this.degreesOfFreedom,
                urdfPath = this.urdfPath,
                srdfPath = this.srdfPath,
                planningGroup = this.planningGroup,
                jointNames = GetJointNames(),
                jointLimitsMin = GetJointLimitsMin(),
                jointLimitsMax = GetJointLimitsMax()
            };
        }

        private string[] GetJointNames()
        {
            string[] names = new string[jointData.Length];
            for (int i = 0; i < jointData.Length; i++)
            {
                names[i] = jointData[i].name;
            }
            return names;
        }

        private float[] GetJointLimitsMin()
        {
            float[] limits = new float[jointData.Length];
            for (int i = 0; i < jointData.Length; i++)
            {
                limits[i] = jointData[i].lowerLimit;
            }
            return limits;
        }

        private float[] GetJointLimitsMax()
        {
            float[] limits = new float[jointData.Length];
            for (int i = 0; i < jointData.Length; i++)
            {
                limits[i] = jointData[i].upperLimit;
            }
            return limits;
        }

        #endregion
    }

    #region Data Structures

    /// <summary>
    /// 关节数据 / Joint Data
    /// </summary>
    [Serializable]
    public class JointData
    {
        public int index;
        public string name;
        public ArticulationJointType jointType;
        public float lowerLimit;
        public float upperLimit;
        public float maxVelocity;
        public float maxAcceleration;
        public float currentPosition;
        public float currentVelocity;
        public float targetPosition;
    }

    /// <summary>
    /// 机器人配置 / Robot Configuration
    /// </summary>
    [Serializable]
    public class RobotConfiguration
    {
        public string robotName;
        public int dof;
        public string urdfPath;
        public string srdfPath;
        public string planningGroup;
        public string[] jointNames;
        public float[] jointLimitsMin;
        public float[] jointLimitsMax;
    }

    #endregion
}