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

using System;
using UnityEngine;

namespace UnityMoveIt2.Core
{
    /// <summary>
    /// 交互式末端执行器控制器 / Interactive End Effector Controller
    /// 实现6DOF拖拽控制功能 / Implements 6DOF drag control functionality
    /// </summary>
    public class InteractiveEndEffectorController : MonoBehaviour
    {
        #region Inspector Fields

        [Header("控制设置 / Control Settings")]
        [SerializeField]
        [Tooltip("启用位置拖拽 / Enable Position Drag")]
        private bool enablePositionDrag = true;

        [SerializeField]
        [Tooltip("启用旋转拖拽 / Enable Rotation Drag")]
        private bool enableRotationDrag = true;

        [SerializeField]
        [Tooltip("拖拽灵敏度 / Drag Sensitivity")]
        private float dragSensitivity = 1.0f;

        [SerializeField]
        [Tooltip("旋转灵敏度 / Rotation Sensitivity")]
        private float rotationSensitivity = 1.0f;

        [Header("约束设置 / Constraint Settings")]
        [SerializeField]
        [Tooltip("遵守关节限制 / Respect Joint Limits")]
        private bool respectJointLimits = true;

        [SerializeField]
        [Tooltip("避免碰撞 / Avoid Collisions")]
        private bool avoidCollisions = true;

        [SerializeField]
        [Tooltip("工作空间半径 / Workspace Radius")]
        private float workspaceRadius = 2.0f;

        [SerializeField]
        [Tooltip("吸附到网格 / Snap to Grid")]
        private bool snapToGrid = false;

        [SerializeField]
        [Tooltip("网格大小 / Grid Size")]
        private float gridSize = 0.01f;

        [Header("视觉反馈 / Visual Feedback")]
        [SerializeField]
        [Tooltip("显示目标位姿 / Show Target Pose")]
        private bool showTargetPose = true;

        [SerializeField]
        [Tooltip("目标位姿颜色 / Target Pose Color")]
        private Color targetPoseColor = Color.cyan;

        [SerializeField]
        [Tooltip("无效位姿颜色 / Invalid Pose Color")]
        private Color invalidPoseColor = Color.red;

        [Header("引用 / References")]
        [SerializeField]
        private UniversalRobotAdapter robotAdapter;

        [SerializeField]
        private Camera mainCamera;

        #endregion

        #region Private Fields

        private Vector3 targetPosition;
        private Quaternion targetRotation;
        private bool isDragging = false;
        private bool isRotating = false;
        private Vector3 dragStartPosition;
        private Quaternion dragStartRotation;
        private Vector3 mouseStartPosition;
        private Plane dragPlane;
        private GameObject targetVisualizer;
        private MeshRenderer targetRenderer;
        private bool isTargetValid = true;
        private float lastIKRequestTime;
        private const float IK_REQUEST_INTERVAL = 0.1f; // 100ms between IK requests

        #endregion

        #region Properties

        public Vector3 TargetPosition => targetPosition;
        public Quaternion TargetRotation => targetRotation;
        public bool IsDragging => isDragging;
        public bool IsTargetValid => isTargetValid;

        #endregion

        #region Events

        public event Action<Vector3, Quaternion> OnTargetPoseChanged;
        public event Action<float[]> OnIKSolutionReceived;
        public event Action<string> OnIKFailed;

        #endregion

        #region Unity Lifecycle

        private void Awake()
        {
            // 自动获取组件引用 / Auto-get component references
            if (robotAdapter == null)
            {
                robotAdapter = GetComponentInParent<UniversalRobotAdapter>();
            }

            if (mainCamera == null)
            {
                mainCamera = Camera.main;
            }

            // 创建目标可视化器 / Create target visualizer
            CreateTargetVisualizer();
        }

        private void Start()
        {
            if (robotAdapter == null)
            {
                Debug.LogError("[InteractiveEndEffectorController] 未找到UniversalRobotAdapter组件 / UniversalRobotAdapter not found");
                enabled = false;
                return;
            }

            // 初始化目标位姿为当前末端执行器位姿 / Initialize target pose to current end effector pose
            Pose currentPose = robotAdapter.GetEndEffectorPose();
            targetPosition = currentPose.position;
            targetRotation = currentPose.rotation;

            UpdateTargetVisualizer();
        }

        private void Update()
        {
            HandleInput();
            UpdateTargetVisualizer();
        }

        #endregion

        #region Input Handling

        private void HandleInput()
        {
            // 左键拖拽位置 / Left click for position drag
            if (enablePositionDrag)
            {
                if (Input.GetMouseButtonDown(0) && !isDragging && !isRotating)
                {
                    StartPositionDrag();
                }
                else if (Input.GetMouseButton(0) && isDragging)
                {
                    UpdatePositionDrag();
                }
                else if (Input.GetMouseButtonUp(0) && isDragging)
                {
                    EndPositionDrag();
                }
            }

            // 右键拖拽旋转 / Right click for rotation drag
            if (enableRotationDrag)
            {
                if (Input.GetMouseButtonDown(1) && !isDragging && !isRotating)
                {
                    StartRotationDrag();
                }
                else if (Input.GetMouseButton(1) && isRotating)
                {
                    UpdateRotationDrag();
                }
                else if (Input.GetMouseButtonUp(1) && isRotating)
                {
                    EndRotationDrag();
                }
            }

            // 键盘快捷键 / Keyboard shortcuts
            HandleKeyboardInput();
        }

        private void HandleKeyboardInput()
        {
            // 按住Shift增加精度 / Hold Shift for increased precision
            float precision = Input.GetKey(KeyCode.LeftShift) ? 0.1f : 1.0f;

            // 方向键微调位置 / Arrow keys for fine position adjustment
            if (Input.GetKey(KeyCode.UpArrow))
            {
                targetPosition += Vector3.forward * 0.01f * precision;
                RequestIKSolution();
            }
            if (Input.GetKey(KeyCode.DownArrow))
            {
                targetPosition += Vector3.back * 0.01f * precision;
                RequestIKSolution();
            }
            if (Input.GetKey(KeyCode.LeftArrow))
            {
                targetPosition += Vector3.left * 0.01f * precision;
                RequestIKSolution();
            }
            if (Input.GetKey(KeyCode.RightArrow))
            {
                targetPosition += Vector3.right * 0.01f * precision;
                RequestIKSolution();
            }
            if (Input.GetKey(KeyCode.PageUp))
            {
                targetPosition += Vector3.up * 0.01f * precision;
                RequestIKSolution();
            }
            if (Input.GetKey(KeyCode.PageDown))
            {
                targetPosition += Vector3.down * 0.01f * precision;
                RequestIKSolution();
            }

            // R键重置到当前位姿 / R key to reset to current pose
            if (Input.GetKeyDown(KeyCode.R))
            {
                ResetToCurrentPose();
            }
        }

        #endregion

        #region Position Drag

        private void StartPositionDrag()
        {
            Ray ray = mainCamera.ScreenPointToRay(Input.mousePosition);

            // 检查是否点击到末端执行器附近 / Check if clicked near end effector
            if (!IsNearEndEffector(ray.origin, ray.direction))
            {
                return;
            }

            isDragging = true;
            dragStartPosition = targetPosition;
            mouseStartPosition = Input.mousePosition;

            // 创建拖拽平面(垂直于相机) / Create drag plane perpendicular to camera
            Vector3 cameraForward = mainCamera.transform.forward;
            dragPlane = new Plane(cameraForward, targetPosition);

            Debug.Log("[InteractiveEndEffectorController] 开始位置拖拽 / Started position drag");
        }

        private void UpdatePositionDrag()
        {
            Ray ray = mainCamera.ScreenPointToRay(Input.mousePosition);
            float enter;

            if (dragPlane.Raycast(ray, out enter))
            {
                Vector3 hitPoint = ray.GetPoint(enter);
                Vector3 offset = hitPoint - dragStartPosition;

                // 应用灵敏度 / Apply sensitivity
                targetPosition = dragStartPosition + offset * dragSensitivity;

                // 应用约束 / Apply constraints
                ApplyConstraints();

                // 请求IK解算 / Request IK solution
                if (Time.time - lastIKRequestTime > IK_REQUEST_INTERVAL)
                {
                    RequestIKSolution();
                    lastIKRequestTime = Time.time;
                }
            }
        }

        private void EndPositionDrag()
        {
            isDragging = false;
            Debug.Log($"[InteractiveEndEffectorController] 结束位置拖拽，目标位置: {targetPosition}");

            // 最终IK求解 / Final IK solution
            RequestIKSolution();
        }

        #endregion

        #region Rotation Drag

        private void StartRotationDrag()
        {
            Ray ray = mainCamera.ScreenPointToRay(Input.mousePosition);

            if (!IsNearEndEffector(ray.origin, ray.direction))
            {
                return;
            }

            isRotating = true;
            dragStartRotation = targetRotation;
            mouseStartPosition = Input.mousePosition;

            Debug.Log("[InteractiveEndEffectorController] 开始旋转拖拽 / Started rotation drag");
        }

        private void UpdateRotationDrag()
        {
            Vector2 mouseDelta = (Vector2)Input.mousePosition - (Vector2)mouseStartPosition;

            // 水平移动控制Yaw / Horizontal movement controls yaw
            float yaw = mouseDelta.x * rotationSensitivity * 0.5f;

            // 垂直移动控制Pitch / Vertical movement controls pitch
            float pitch = mouseDelta.y * rotationSensitivity * 0.5f;

            // 按住Ctrl控制Roll / Hold Ctrl for roll
            float roll = Input.GetKey(KeyCode.LeftControl) ? mouseDelta.x * rotationSensitivity * 0.5f : 0f;

            // 应用旋转 / Apply rotation
            Quaternion deltaRotation = Quaternion.Euler(pitch, yaw, roll);
            targetRotation = dragStartRotation * deltaRotation;

            // 请求IK解算 / Request IK solution
            if (Time.time - lastIKRequestTime > IK_REQUEST_INTERVAL)
            {
                RequestIKSolution();
                lastIKRequestTime = Time.time;
            }
        }

        private void EndRotationDrag()
        {
            isRotating = false;
            Debug.Log($"[InteractiveEndEffectorController] 结束旋转拖拽，目标旋转: {targetRotation.eulerAngles}");

            // 最终IK求解 / Final IK solution
            RequestIKSolution();
        }

        #endregion

        #region Constraints

        private void ApplyConstraints()
        {
            // 工作空间约束 / Workspace constraint
            Vector3 basePosition = robotAdapter.BaseLink.position;
            Vector3 offset = targetPosition - basePosition;

            if (offset.magnitude > workspaceRadius)
            {
                targetPosition = basePosition + offset.normalized * workspaceRadius;
            }

            // 网格吸附 / Grid snapping
            if (snapToGrid)
            {
                targetPosition.x = Mathf.Round(targetPosition.x / gridSize) * gridSize;
                targetPosition.y = Mathf.Round(targetPosition.y / gridSize) * gridSize;
                targetPosition.z = Mathf.Round(targetPosition.z / gridSize) * gridSize;
            }
        }

        #endregion

        #region IK Solution

        private void RequestIKSolution()
        {
            // 触发事件通知其他组件请求IK解算 / Trigger event to notify other components
            OnTargetPoseChanged?.Invoke(targetPosition, targetRotation);

            // 这里应该通过ROS通信层请求IK解算 / Here should request IK solution through ROS communication layer
            // 暂时假设目标有效 / Temporarily assume target is valid
            isTargetValid = true;
        }

        /// <summary>
        /// 接收IK解算结果 / Receive IK Solution Result
        /// 由通信层调用 / Called by communication layer
        /// </summary>
        public void ReceiveIKSolution(float[] jointPositions, bool success)
        {
            if (success && jointPositions != null)
            {
                isTargetValid = true;
                OnIKSolutionReceived?.Invoke(jointPositions);

                Debug.Log($"[InteractiveEndEffectorController] IK解算成功 / IK solution succeeded");
            }
            else
            {
                isTargetValid = false;
                OnIKFailed?.Invoke("IK解算失败 / IK solution failed");

                Debug.LogWarning("[InteractiveEndEffectorController] IK解算失败 / IK solution failed");
            }
        }

        #endregion

        #region Visualization

        private void CreateTargetVisualizer()
        {
            if (!showTargetPose)
                return;

            // 创建目标可视化GameObject / Create target visualizer GameObject
            targetVisualizer = GameObject.CreatePrimitive(PrimitiveType.Sphere);
            targetVisualizer.name = "TargetPoseVisualizer";
            targetVisualizer.transform.SetParent(transform);
            targetVisualizer.transform.localScale = Vector3.one * 0.05f;

            // 设置材质 / Setup material
            targetRenderer = targetVisualizer.GetComponent<MeshRenderer>();
            Material mat = new Material(Shader.Find("Standard"));
            mat.color = targetPoseColor;
            mat.SetFloat("_Mode", 3); // Transparent mode
            mat.SetInt("_SrcBlend", (int)UnityEngine.Rendering.BlendMode.SrcAlpha);
            mat.SetInt("_DstBlend", (int)UnityEngine.Rendering.BlendMode.OneMinusSrcAlpha);
            mat.SetInt("_ZWrite", 0);
            mat.DisableKeyword("_ALPHATEST_ON");
            mat.EnableKeyword("_ALPHABLEND_ON");
            mat.DisableKeyword("_ALPHAPREMULTIPLY_ON");
            mat.renderQueue = 3000;
            mat.color = new Color(targetPoseColor.r, targetPoseColor.g, targetPoseColor.b, 0.5f);
            targetRenderer.material = mat;

            // 禁用碰撞 / Disable collision
            Destroy(targetVisualizer.GetComponent<Collider>());

            // 添加坐标轴指示 / Add coordinate axes indicator
            CreateCoordinateAxes(targetVisualizer.transform);
        }

        private void CreateCoordinateAxes(Transform parent)
        {
            float axisLength = 0.1f;
            float axisWidth = 0.005f;

            // X轴 - 红色 / X-axis - Red
            CreateAxis(parent, Vector3.right, Color.red, axisLength, axisWidth);

            // Y轴 - 绿色 / Y-axis - Green
            CreateAxis(parent, Vector3.up, Color.green, axisLength, axisWidth);

            // Z轴 - 蓝色 / Z-axis - Blue
            CreateAxis(parent, Vector3.forward, Color.blue, axisLength, axisWidth);
        }

        private void CreateAxis(Transform parent, Vector3 direction, Color color, float length, float width)
        {
            GameObject axis = GameObject.CreatePrimitive(PrimitiveType.Cylinder);
            axis.transform.SetParent(parent);
            axis.transform.localPosition = direction * length * 0.5f;
            axis.transform.localRotation = Quaternion.FromToRotation(Vector3.up, direction);
            axis.transform.localScale = new Vector3(width, length * 0.5f, width);

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

            Destroy(axis.GetComponent<Collider>());
        }

        private void UpdateTargetVisualizer()
        {
            if (!showTargetPose || targetVisualizer == null)
                return;

            targetVisualizer.transform.position = targetPosition;
            targetVisualizer.transform.rotation = targetRotation;

            // 根据有效性更新颜色 / Update color based on validity
            if (targetRenderer != null)
            {
                Color color = isTargetValid ? targetPoseColor : invalidPoseColor;
                Material mat = targetRenderer.material;
                mat.color = new Color(color.r, color.g, color.b, 0.5f);
            }

            // 仅在拖拽时显示 / Only show when dragging
            targetVisualizer.SetActive(isDragging || isRotating);
        }

        #endregion

        #region Utility Methods

        private bool IsNearEndEffector(Vector3 rayOrigin, Vector3 rayDirection)
        {
            if (robotAdapter.EndEffector == null)
                return false;

            Vector3 endEffectorPos = robotAdapter.EndEffector.position;
            Vector3 closestPoint = rayOrigin + Vector3.Project(endEffectorPos - rayOrigin, rayDirection);
            float distance = Vector3.Distance(closestPoint, endEffectorPos);

            // 判断是否在末端执行器附近(50cm范围内) / Check if near end effector (within 50cm)
            return distance < 0.5f;
        }

        private void ResetToCurrentPose()
        {
            Pose currentPose = robotAdapter.GetEndEffectorPose();
            targetPosition = currentPose.position;
            targetRotation = currentPose.rotation;

            Debug.Log("[InteractiveEndEffectorController] 重置到当前位姿 / Reset to current pose");
        }

        #endregion

        #region Public Methods

        /// <summary>
        /// 设置目标位姿 / Set Target Pose
        /// </summary>
        public void SetTargetPose(Vector3 position, Quaternion rotation)
        {
            targetPosition = position;
            targetRotation = rotation;
            RequestIKSolution();
        }

        /// <summary>
        /// 设置目标位置 / Set Target Position
        /// </summary>
        public void SetTargetPosition(Vector3 position)
        {
            targetPosition = position;
            RequestIKSolution();
        }

        /// <summary>
        /// 设置目标旋转 / Set Target Rotation
        /// </summary>
        public void SetTargetRotation(Quaternion rotation)
        {
            targetRotation = rotation;
            RequestIKSolution();
        }

        #endregion

        #region Gizmos

        private void OnDrawGizmos()
        {
            if (!Application.isPlaying || robotAdapter == null)
                return;

            // 绘制工作空间边界 / Draw workspace boundary
            Gizmos.color = Color.yellow;
            Gizmos.DrawWireSphere(robotAdapter.BaseLink.position, workspaceRadius);

            // 绘制从基座到目标的连线 / Draw line from base to target
            if (isDragging || isRotating)
            {
                Gizmos.color = isTargetValid ? Color.green : Color.red;
                Gizmos.DrawLine(robotAdapter.BaseLink.position, targetPosition);
            }
        }

        #endregion
    }
}