unity2moveit2/unity-project/Assets/Scripts/Core/VisualizationEngine.JointState.cs

using System.Collections.Generic;
using UnityEngine;

namespace UnityMoveIt2.Core
{
    /// <summary>
    /// 可视化引擎 - 关节状态可视化部分 / Visualization Engine - Joint State Visualization Part
    /// </summary>
    public partial class VisualizationEngine
    {
        #region Joint State Fields

        [Header("关节状态可视化设置 / Joint State Visualization Settings")]
        [SerializeField]
        [Tooltip("显示关节轴 / Show Joint Axes")]
        private bool showJointAxes = false;

        [SerializeField]
        [Tooltip("轴长度 / Axis Length")]
        private float axisLength = 0.1f;

        [SerializeField]
        [Tooltip("显示关节限制 / Show Joint Limits")]
        private bool showJointLimits = true;

        [SerializeField]
        private Color limitColor = Color.magenta;

        [SerializeField]
        [Tooltip("显示关节名称 / Show Joint Names")]
        private bool showJointNames = false;

        private List<GameObject> jointAxisMarkers;
        private List<GameObject> jointLimitIndicators;
        private List<GameObject> jointNameLabels;

        #endregion

        #region Joint State Initialization

        private void InitializeJointStateVisualization()
        {
            jointAxisMarkers = new List<GameObject>();
            jointLimitIndicators = new List<GameObject>();
            jointNameLabels = new List<GameObject>();

            Debug.Log("[VisualizationEngine] 关节状态可视化已初始化 / Joint state visualization initialized");
        }

        #endregion

        #region Joint State Update

        private void UpdateJointStateVisualization(float[] jointPositions)
        {
            if (robotAdapter == null || jointPositions == null)
                return;

            // 更新关节轴显示 / Update joint axes display
            if (showJointAxes)
            {
                UpdateJointAxes();
            }

            // 更新关节限制指示器 / Update joint limit indicators
            if (showJointLimits)
            {
                UpdateJointLimits(jointPositions);
            }

            // 更新关节名称标签 / Update joint name labels
            if (showJointNames)
            {
                UpdateJointNames();
            }
        }

        private void UpdateJointAxes()
        {
            if (robotAdapter == null || robotAdapter.JointData == null)
                return;

            // 确保标记数量匹配关节数量 / Ensure marker count matches joint count
            while (jointAxisMarkers.Count < robotAdapter.DOF)
            {
                CreateJointAxisMarker(jointAxisMarkers.Count);
            }

            // 更新标记位置和方向 / Update marker positions and orientations
            for (int i = 0; i < robotAdapter.DOF && i < jointAxisMarkers.Count; i++)
            {
                if (jointAxisMarkers[i] != null)
                {
                    Pose linkPose = robotAdapter.GetLinkPose(i);
                    jointAxisMarkers[i].transform.position = linkPose.position;
                    jointAxisMarkers[i].transform.rotation = linkPose.rotation;
                }
            }
        }

        private void CreateJointAxisMarker(int jointIndex)
        {
            GameObject marker = new GameObject($"JointAxis_{jointIndex}");
            marker.transform.SetParent(jointStateContainer.transform);

            // 创建三个坐标轴 / Create three coordinate axes
            CreateAxisLine(marker.transform, Vector3.right, Color.red, axisLength);    // X轴
            CreateAxisLine(marker.transform, Vector3.up, Color.green, axisLength);      // Y轴
            CreateAxisLine(marker.transform, Vector3.forward, Color.blue, axisLength);  // Z轴

            jointAxisMarkers.Add(marker);
        }

        private void CreateAxisLine(Transform parent, Vector3 direction, Color color, float length)
        {
            GameObject axis = new GameObject($"Axis_{direction}");
            axis.transform.SetParent(parent);
            axis.transform.localPosition = Vector3.zero;

            LineRenderer line = axis.AddComponent<LineRenderer>();
            line.startWidth = 0.005f;
            line.endWidth = 0.005f;
            line.material = new Material(Shader.Find("Sprites/Default"));
            line.startColor = color;
            line.endColor = color;
            line.positionCount = 2;
            line.useWorldSpace = false;
            line.SetPosition(0, Vector3.zero);
            line.SetPosition(1, direction * length);
        }

        private void UpdateJointLimits(float[] jointPositions)
        {
            if (robotAdapter == null || robotAdapter.JointData == null)
                return;

            // 确保指示器数量匹配 / Ensure indicator count matches
            while (jointLimitIndicators.Count < robotAdapter.DOF)
            {
                CreateJointLimitIndicator(jointLimitIndicators.Count);
            }

            // 更新限制指示器 / Update limit indicators
            for (int i = 0; i < robotAdapter.DOF && i < jointPositions.Length; i++)
            {
                if (jointLimitIndicators[i] != null)
                {
                    UpdateLimitIndicatorColor(i, jointPositions[i]);
                }
            }
        }

        private void CreateJointLimitIndicator(int jointIndex)
        {
            GameObject indicator = GameObject.CreatePrimitive(PrimitiveType.Sphere);
            indicator.name = $"JointLimitIndicator_{jointIndex}";
            indicator.transform.SetParent(jointStateContainer.transform);
            indicator.transform.localScale = Vector3.one * 0.02f;

            // 移除碰撞器 / Remove collider
            Collider collider = indicator.GetComponent<Collider>();
            if (collider != null)
            {
                Destroy(collider);
            }

            MeshRenderer renderer = indicator.GetComponent<MeshRenderer>();
            Material mat = new Material(Shader.Find("Standard"));
            mat.color = Color.green;
            renderer.material = mat;

            jointLimitIndicators.Add(indicator);
        }

        private void UpdateLimitIndicatorColor(int jointIndex, float currentPosition)
        {
            if (robotAdapter.JointData == null || jointIndex >= robotAdapter.JointData.Length)
                return;

            JointData joint = robotAdapter.JointData[jointIndex];
            float range = joint.upperLimit - joint.lowerLimit;
            float normalizedPos = (currentPosition - joint.lowerLimit) / range;

            // 计算颜色 / Calculate color
            Color color;
            if (normalizedPos < 0.1f || normalizedPos > 0.9f)
            {
                color = limitColor; // 接近限制 / Near limit
            }
            else if (normalizedPos < 0.2f || normalizedPos > 0.8f)
            {
                color = Color.yellow; // 警告区 / Warning zone
            }
            else
            {
                color = Color.green; // 安全区 / Safe zone
            }

            MeshRenderer renderer = jointLimitIndicators[jointIndex].GetComponent<MeshRenderer>();
            if (renderer != null)
            {
                renderer.material.color = color;
            }

            // 更新位置 / Update position
            Pose linkPose = robotAdapter.GetLinkPose(jointIndex);
            jointLimitIndicators[jointIndex].transform.position = linkPose.position;
        }

        private void UpdateJointNames()
        {
            if (robotAdapter == null || robotAdapter.JointData == null)
                return;

            // 确保标签数量匹配 / Ensure label count matches
            while (jointNameLabels.Count < robotAdapter.DOF)
            {
                CreateJointNameLabel(jointNameLabels.Count);
            }

            // 更新标签位置 / Update label positions
            for (int i = 0; i < robotAdapter.DOF && i < jointNameLabels.Count; i++)
            {
                if (jointNameLabels[i] != null)
                {
                    Pose linkPose = robotAdapter.GetLinkPose(i);
                    jointNameLabels[i].transform.position = linkPose.position + Vector3.up * 0.1f;

                    // 让标签始终面向相机 / Make label always face camera
                    Camera mainCam = Camera.main;
                    if (mainCam != null)
                    {
                        jointNameLabels[i].transform.rotation = Quaternion.LookRotation(
                            jointNameLabels[i].transform.position - mainCam.transform.position
                        );
                    }
                }
            }
        }

        private void CreateJointNameLabel(int jointIndex)
        {
            GameObject label = new GameObject($"JointNameLabel_{jointIndex}");
            label.transform.SetParent(jointStateContainer.transform);

            // 注意: 实际项目中应使用TextMeshPro
            // 这里简化处理,仅创建占位对象
            // Note: In actual project, should use TextMeshPro
            // Simplified here, just create placeholder

            jointNameLabels.Add(label);
        }

        #endregion

        #region Public Joint State Methods

        /// <summary>
        /// 清除关节状态可视化 / Clear Joint State Visualization
        /// </summary>
        public void ClearJointStateVisualization()
        {
            foreach (var marker in jointAxisMarkers)
            {
                if (marker != null)
                {
                    Destroy(marker);
                }
            }
            jointAxisMarkers.Clear();

            foreach (var indicator in jointLimitIndicators)
            {
                if (indicator != null)
                {
                    Destroy(indicator);
                }
            }
            jointLimitIndicators.Clear();

            foreach (var label in jointNameLabels)
            {
                if (label != null)
                {
                    Destroy(label);
                }
            }
            jointNameLabels.Clear();

            Debug.Log("[VisualizationEngine] 已清除关节状态可视化 / Cleared joint state visualization");
        }

        /// <summary>
        /// 设置关节轴显示 / Set Joint Axes Display
        /// </summary>
        public void SetShowJointAxes(bool show)
        {
            showJointAxes = show;

            foreach (var marker in jointAxisMarkers)
            {
                if (marker != null)
                {
                    marker.SetActive(show);
                }
            }
        }

        /// <summary>
        /// 设置关节限制显示 / Set Joint Limits Display
        /// </summary>
        public void SetShowJointLimits(bool show)
        {
            showJointLimits = show;

            foreach (var indicator in jointLimitIndicators)
            {
                if (indicator != null)
                {
                    indicator.SetActive(show);
                }
            }
        }

        #endregion
    }
}