ApexSurgery/Assets/Imstk/Scripts/Editor/GraspingManagerEditor.cs

using ImstkUnity;
using System.Collections.Generic;
using System.Linq;
using UnityEditor;
using UnityEngine;

namespace ImstkEditor
{
    [CustomEditor(typeof(GraspingManager))]
    public class GraspingManagerEditor : Editor
    {
        bool _folded = true;
        Rigid _rigid;
        GeometryFilter _graspingGeometry;
        List<GraspingManager.GraspingData> _graspingData;

        List<DynamicalModel> _allDeformables = new List<DynamicalModel>();

        GUIContent _warning;

        double _deformableGraspingStiffness;
        double _rigidGraspingStiffness;

        GraspingManagerEditor()
        {
            EditorApplication.hierarchyChanged += OnHierarchyChanged;
        }

        private void OnHierarchyChanged()
        {
            _allDeformables = (Resources.FindObjectsOfTypeAll(typeof(DynamicalModel)) as DynamicalModel[]).ToList<DynamicalModel>();
        }

        public override void OnInspectorGUI()
        {
            if (_allDeformables.Count == 0) OnHierarchyChanged();

            var script = target as GraspingManager;
            EditorGUI.BeginChangeCheck();
            EditorGUILayout.BeginHorizontal();
            _rigid = EditorGUILayout.ObjectField("Grasping Object", 
                script.rigid, typeof(ImstkUnity.Rigid), true) as ImstkUnity.Rigid;
            if (script.rigid == null)
            {
                _warning = EditorGUIUtility.IconContent("Error", "|Grasping Object Can't be null");
                EditorGUILayout.LabelField(_warning, GUILayout.Width(_warning.image.width));
            }
            EditorGUILayout.EndHorizontal();
            _graspingGeometry = EditorGUILayout.ObjectField("Grasping Geometry (Optional)", 
                script.graspingGeometry, typeof(GeometryFilter), true) as GeometryFilter;

            GUILayout.BeginVertical(EditorStyles.helpBox);

            _deformableGraspingStiffness = EditorGUILayout.DoubleField("Deformable Grasping Stiffness", script.deformableGraspingStiffness);
            _rigidGraspingStiffness = EditorGUILayout.DoubleField("Rigid Grasping Stiffness", script.rigidGraspingStiffness);

            GUILayout.EndVertical();
            // Maybe this is something we should do for all imstk components 
            // as we can't edit them during runtime anyway
            if (_rigid != null && !Application.isPlaying)
            {
                // Takes up a lot of resources when playing, just don't execute 
                HandleGraspingData(script);
                UpdateScript(script);
            }
            else
            {
                EditorGUI.EndChangeCheck();
            }
        }

        private void UpdateScript(GraspingManager script)
        {
            if (EditorGUI.EndChangeCheck())
            {
                Undo.RegisterCompleteObjectUndo(script, "Change of Parameters");
                script.rigid = _rigid;
                script.graspingGeometry = _graspingGeometry;
                script.graspingData = new List<GraspingManager.GraspingData>(_graspingData);
                script.deformableGraspingStiffness = _deformableGraspingStiffness;
                script.rigidGraspingStiffness = _rigidGraspingStiffness;
            }
        }

        private void HandleGraspingData(GraspingManager script)
        {
            // if there aren't any items to grasp, make sure ...
            if (_allDeformables.Count == 0) OnHierarchyChanged();

            _graspingData = new List<GraspingManager.GraspingData>(script.graspingData);

            foreach (var item in _allDeformables)
            {
                if (item is StaticModel || item == _rigid) continue;

                if (_graspingData.FindIndex(x => x.target == item) < 0)
                {
                    GraspingManager.GraspingData graspingData = new GraspingManager.GraspingData();
                    graspingData.enabled = false;
                    graspingData.type = Grasping.GraspType.Cell;
                    graspingData.target = item;
                    _graspingData.Add(graspingData);
                }
            }

            // Remove stale Objects
            _graspingData.RemoveAll(x => !_allDeformables.Contains(x.target));

            _folded = EditorGUILayout.Foldout(_folded, "Graspable Objects");
            if (_folded)
            {
                EditorGUILayout.BeginVertical();
                for (int i = 0; i < _graspingData.Count; ++i)
                {
                    var item = _graspingData[i];
                    EditorGUILayout.BeginHorizontal();
                    item.enabled = EditorGUILayout.Toggle(item.enabled, GUILayout.Width(20));
                    EditorGUI.BeginDisabledGroup(!item.enabled);
                    EditorGUILayout.LabelField(item.target.name);
                    item.type = (Grasping.GraspType)EditorGUILayout.EnumPopup(item.type);
                    EditorGUI.EndDisabledGroup();
                    EditorGUILayout.EndHorizontal();
                    _graspingData[i] = item;
                }
                EditorGUILayout.EndVertical();
            }
        }
    }
}