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

using System.Collections.Generic;
using UnityEngine;

namespace UnityMoveIt2.Core
{
    /// <summary>
    /// 可视化引擎 - 轨迹可视化部分 / Visualization Engine - Trajectory Visualization Part
    /// </summary>
    public partial class VisualizationEngine
    {
        #region Trajectory Fields

        [Header("轨迹可视化设置 / Trajectory Visualization Settings")]
        [SerializeField]
        private Color trajectoryColor = Color.green;

        [SerializeField]
        private Color executingColor = Color.blue;

        [SerializeField]
        private Color completedColor = Color.gray;

        [SerializeField]
        [Tooltip("轨迹线宽度 / Trajectory Line Width")]
        private float lineWidth = 0.02f;

        [SerializeField]
        [Tooltip("显示路径点 / Show Waypoints")]
        private bool showWaypoints = true;

        [SerializeField]
        [Tooltip("路径点大小 / Waypoint Size")]
        private float waypointSize = 0.03f;

        private LineRenderer trajectoryLineRenderer;
        private List<Vector3> trajectoryPoints;
        private List<GameObject> waypointMarkers;
        private float trajectoryProgress = 0f;

        #endregion

        #region Trajectory Initialization

        private void InitializeTrajectoryVisualization()
        {
            trajectoryPoints = new List<Vector3>();
            waypointMarkers = new List<GameObject>();

            // 创建LineRenderer / Create LineRenderer
            GameObject lineObj = GetOrCreateVisualizationObject("TrajectoryLine", trajectoryContainer.transform);
            trajectoryLineRenderer = lineObj.GetComponent<LineRenderer>();

            if (trajectoryLineRenderer == null)
            {
                trajectoryLineRenderer = lineObj.AddComponent<LineRenderer>();
            }

            // 配置LineRenderer / Configure LineRenderer
            trajectoryLineRenderer.startWidth = lineWidth;
            trajectoryLineRenderer.endWidth = lineWidth;
            trajectoryLineRenderer.material = new Material(Shader.Find("Sprites/Default"));
            trajectoryLineRenderer.startColor = trajectoryColor;
            trajectoryLineRenderer.endColor = trajectoryColor;
            trajectoryLineRenderer.positionCount = 0;
            trajectoryLineRenderer.useWorldSpace = true;

            Debug.Log("[VisualizationEngine] 轨迹可视化已初始化 / Trajectory visualization initialized");
        }

        #endregion

        #region Trajectory Update

        private void UpdateTrajectoryVisualization()
        {
            if (trajectoryLineRenderer == null || trajectoryPoints.Count == 0)
                return;

            // 更新LineRenderer位置 / Update LineRenderer positions
            trajectoryLineRenderer.positionCount = trajectoryPoints.Count;
            trajectoryLineRenderer.SetPositions(trajectoryPoints.ToArray());

            // 更新颜色渐变以显示执行进度 / Update color gradient to show execution progress
            if (trajectoryProgress > 0f)
            {
                UpdateTrajectoryGradient();
            }

            // 更新路径点标记 / Update waypoint markers
            if (showWaypoints)
            {
                UpdateWaypointMarkers();
            }
        }

        private void UpdateTrajectoryGradient()
        {
            Gradient gradient = new Gradient();
            GradientColorKey[] colorKeys = new GradientColorKey[3];
            GradientAlphaKey[] alphaKeys = new GradientAlphaKey[2];

            // 已完成部分 / Completed part
            colorKeys[0] = new GradientColorKey(completedColor, 0.0f);

            // 当前执行点 / Current execution point
            colorKeys[1] = new GradientColorKey(executingColor, trajectoryProgress);

            // 未执行部分 / Remaining part
            colorKeys[2] = new GradientColorKey(trajectoryColor, 1.0f);

            alphaKeys[0] = new GradientAlphaKey(1.0f, 0.0f);
            alphaKeys[1] = new GradientAlphaKey(1.0f, 1.0f);

            gradient.SetKeys(colorKeys, alphaKeys);
            trajectoryLineRenderer.colorGradient = gradient;
        }

        private void UpdateWaypointMarkers()
        {
            // 确保标记数量匹配路径点数量 / Ensure marker count matches waypoint count
            while (waypointMarkers.Count < trajectoryPoints.Count)
            {
                CreateWaypointMarker(waypointMarkers.Count);
            }

            while (waypointMarkers.Count > trajectoryPoints.Count)
            {
                int lastIndex = waypointMarkers.Count - 1;
                if (waypointMarkers[lastIndex] != null)
                {
                    Destroy(waypointMarkers[lastIndex]);
                }
                waypointMarkers.RemoveAt(lastIndex);
            }

            // 更新标记位置 / Update marker positions
            for (int i = 0; i < trajectoryPoints.Count; i++)
            {
                if (waypointMarkers[i] != null)
                {
                    waypointMarkers[i].transform.position = trajectoryPoints[i];
                }
            }
        }

        private void CreateWaypointMarker(int index)
        {
            GameObject marker = GameObject.CreatePrimitive(PrimitiveType.Sphere);
            marker.name = $"Waypoint_{index}";
            marker.transform.SetParent(trajectoryContainer.transform);
            marker.transform.localScale = Vector3.one * waypointSize;

            // 配置材质 / Configure material
            MeshRenderer renderer = marker.GetComponent<MeshRenderer>();
            Material mat = new Material(Shader.Find("Standard"));
            mat.color = trajectoryColor;
            renderer.material = mat;

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

            waypointMarkers.Add(marker);
        }

        #endregion

        #region Public Trajectory Methods

        /// <summary>
        /// 可视化轨迹 / Visualize Trajectory
        /// </summary>
        public void VisualizeTrajectory(List<Pose> waypoints)
        {
            if (waypoints == null || waypoints.Count == 0)
            {
                Debug.LogWarning("[VisualizationEngine] 轨迹为空 / Trajectory is empty");
                return;
            }

            trajectoryPoints.Clear();

            foreach (var waypoint in waypoints)
            {
                trajectoryPoints.Add(waypoint.position);
            }

            // 限制点数 / Limit point count
            if (trajectoryPoints.Count > maxTrajectoryPoints)
            {
                trajectoryPoints = DownsampleTrajectory(trajectoryPoints, maxTrajectoryPoints);
            }

            trajectoryProgress = 0f;
            UpdateTrajectoryVisualization();

            Debug.Log($"[VisualizationEngine] 可视化轨迹: {trajectoryPoints.Count} 个点 / Visualized trajectory: {trajectoryPoints.Count} points");
        }

        /// <summary>
        /// 从关节轨迹生成末端执行器轨迹可视化 / Generate End Effector Trajectory from Joint Trajectory
        /// </summary>
        public void VisualizeJointTrajectory(List<float[]> jointTrajectory)
        {
            if (jointTrajectory == null || jointTrajectory.Count == 0 || robotAdapter == null)
            {
                Debug.LogWarning("[VisualizationEngine] 无法从关节轨迹生成可视化 / Cannot generate visualization from joint trajectory");
                return;
            }

            List<Pose> endEffectorPoses = new List<Pose>();

            // 保存当前关节状态 / Save current joint state
            float[] originalPositions = robotAdapter.GetCurrentJointPositions();

            // 对每个关节状态计算末端执行器位姿 / Calculate end effector pose for each joint state
            foreach (var jointPositions in jointTrajectory)
            {
                robotAdapter.SetJointPositions(jointPositions);
                Pose endEffectorPose = robotAdapter.GetEndEffectorPose();
                endEffectorPoses.Add(endEffectorPose);
            }

            // 恢复原始关节状态 / Restore original joint state
            robotAdapter.SetJointPositions(originalPositions);

            // 可视化 / Visualize
            VisualizeTrajectory(endEffectorPoses);
        }

        /// <summary>
        /// 更新轨迹执行进度 / Update Trajectory Execution Progress
        /// </summary>
        public void UpdateTrajectoryProgress(float progress)
        {
            trajectoryProgress = Mathf.Clamp01(progress);
            UpdateTrajectoryVisualization();
        }

        /// <summary>
        /// 清除轨迹可视化 / Clear Trajectory Visualization
        /// </summary>
        public void ClearTrajectoryVisualization()
        {
            trajectoryPoints.Clear();

            if (trajectoryLineRenderer != null)
            {
                trajectoryLineRenderer.positionCount = 0;
            }

            foreach (var marker in waypointMarkers)
            {
                if (marker != null)
                {
                    Destroy(marker);
                }
            }
            waypointMarkers.Clear();

            trajectoryProgress = 0f;

            Debug.Log("[VisualizationEngine] 已清除轨迹可视化 / Cleared trajectory visualization");
        }

        #endregion

        #region Trajectory Utilities

        /// <summary>
        /// 下采样轨迹点 / Downsample Trajectory Points
        /// </summary>
        private List<Vector3> DownsampleTrajectory(List<Vector3> points, int targetCount)
        {
            if (points.Count <= targetCount)
                return points;

            List<Vector3> downsampled = new List<Vector3>();
            float step = (float)(points.Count - 1) / (targetCount - 1);

            for (int i = 0; i < targetCount; i++)
            {
                int index = Mathf.RoundToInt(i * step);
                downsampled.Add(points[index]);
            }

            return downsampled;
        }

        /// <summary>
        /// 计算轨迹总长度 / Calculate Total Trajectory Length
        /// </summary>
        public float GetTrajectoryLength()
        {
            if (trajectoryPoints.Count < 2)
                return 0f;

            float totalLength = 0f;
            for (int i = 1; i < trajectoryPoints.Count; i++)
            {
                totalLength += Vector3.Distance(trajectoryPoints[i - 1], trajectoryPoints[i]);
            }

            return totalLength;
        }

        /// <summary>
        /// 获取轨迹上指定进度的位置 / Get Position at Specified Progress on Trajectory
        /// </summary>
        public Vector3 GetTrajectoryPositionAtProgress(float progress)
        {
            progress = Mathf.Clamp01(progress);

            if (trajectoryPoints.Count == 0)
                return Vector3.zero;

            if (trajectoryPoints.Count == 1)
                return trajectoryPoints[0];

            float totalLength = GetTrajectoryLength();
            float targetLength = totalLength * progress;
            float currentLength = 0f;

            for (int i = 1; i < trajectoryPoints.Count; i++)
            {
                float segmentLength = Vector3.Distance(trajectoryPoints[i - 1], trajectoryPoints[i]);
                if (currentLength + segmentLength >= targetLength)
                {
                    float t = (targetLength - currentLength) / segmentLength;
                    return Vector3.Lerp(trajectoryPoints[i - 1], trajectoryPoints[i], t);
                }
                currentLength += segmentLength;
            }

            return trajectoryPoints[trajectoryPoints.Count - 1];
        }

        #endregion
    }
}