NanoBrain-unitypackage/Assets/NanoBrain/Editor/NanoBrain_Editor.cs

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

[CustomEditor(typeof(NanoBrain))]
public class NanoBrain_Editor : Editor {
    private Nucleus currentNucleus;
    private List<NeuroidLayer> layers = new();
    private Dictionary<Nucleus, Vector2Int> neuroidPositions = new();

    protected bool breakOnWake = false;

    #region Start

    private void OnEnable() {
        // NanoBrain brain = target as NanoBrain;
        // if (brain == null)
        //     return;

        // if (brain.rootNucleus == null) {
        //     brain.rootNucleus = CreateInstance<NucleusObj>();
        //     EditorUtility.SetDirty(brain.rootNucleus);
        // }

        SelectNeuron();
    }

    private void SelectNeuron() {
        // GameObject selectedObject = ((NanoBrain)target).gameObject;
        // if (!selectedObject.TryGetComponent(out Boid boid))
        //     return;

        // Nucleus neuroid = boid.totalForce;
        // this.currentNucleus = neuroid;

        // BuildLayers();
        // Debug.Log($"Layercount = {this.layers.Count}");
    }

    #endregion Start    

    #region Update

    public override void OnInspectorGUI() {
        if (this.currentNucleus == null)
            return;

        breakOnWake = EditorGUILayout.Toggle("Break on wake", breakOnWake);
        if (breakOnWake && currentNucleus is Neuroid currentNeuroid) {
            if (!currentNeuroid.isSleeping)
                Debug.Break();
        }

        DrawGraph();

        EditorGUILayout.TextField("Name", currentNucleus.name);
        EditorGUILayout.Vector3Field("Output Value", currentNucleus.outputValue.ToVector3());
        if (currentNucleus.synapses.Count > 0) {
            EditorGUI.indentLevel++;
            //foreach ((Nucleus nucleus, float weight) in currentNucleus.synapses) {
            foreach (Synapse synapse in currentNucleus.synapses) {
                Nucleus nucleus = synapse.nucleus;
                float weight = synapse.weight;

                EditorGUI.BeginDisabledGroup(nucleus.isSleeping);

                EditorGUILayout.BeginHorizontal();
                EditorGUILayout.Vector3Field(nucleus.name, nucleus.outputValue.ToVector3());
                EditorGUILayout.FloatField(weight, GUILayout.Width(50));
                EditorGUILayout.EndHorizontal();

                EditorGUI.EndDisabledGroup();

            }
            EditorGUI.indentLevel--;
        }
        DrawDefaultInspector();
    }

    private void BuildLayers() {
        // A temporary list to track what's been added to layers
        this.layers = new();
        int layerIx = 0;

        Nucleus selectedNucleus = this.currentNucleus;
        if (selectedNucleus == null)
            return;
        NeuroidLayer currentLayer = new() { ix = layerIx };

        //foreach (Neuroid outputNeuroid in selectedNucleus.receivers) {
        foreach (Receiver receiver in selectedNucleus.receivers) {
            Nucleus outputNeuroid = receiver.nucleus;    
            if (outputNeuroid != null) {
                AddToLayer(currentLayer, outputNeuroid);
                // Debug.Log($"layer {layerIx} nucleus {outputNeuroid.name}");
            }
        }
        if (currentLayer.neuroids.Count > 0) {
            this.layers.Add(currentLayer);
            layerIx++;
            currentLayer = new() { ix = layerIx };
        }

        AddToLayer(currentLayer, selectedNucleus);
        this.layers.Add(currentLayer);
        // Debug.Log($"layer {layerIx} nucleus {selectedNucleus.name}");

        layerIx++;
        currentLayer = new() { ix = layerIx };

        //foreach (Nucleus input in selectedNucleus.synapses.Keys) {
        foreach (Synapse synapse in selectedNucleus.synapses) {
            Nucleus input = synapse.nucleus;
            AddToLayer(currentLayer, input);
            // Debug.Log($"layer {layerIx} nucleus {input.name}");
        }
        if (currentLayer.neuroids.Count > 0) {
            this.layers.Add(currentLayer);
        }
    }

    private void AddToLayer(NeuroidLayer layer, Nucleus nucleus) {
        layer.neuroids.Add(nucleus);
        nucleus.layerIx = layer.ix;
        // Store its position
        Vector2Int neuroidPosition = new(layer.ix, layer.neuroids.Count - 1);
        neuroidPositions[nucleus] = neuroidPosition;

    }

    private void DrawGraph() {
        if (currentNucleus == null)
            return;
        Rect outer = EditorGUILayout.GetControlRect(false, 420);
        GUI.BeginGroup(outer);
        foreach (NeuroidLayer layer in layers)
            DrawLayer(layer);
        GUI.EndGroup();
    }

    private void DrawLayer(NeuroidLayer layer) {
        int nodeCount = layer.neuroids.Count;
        float maxValue = 0;
        foreach (Nucleus nucleus in layer.neuroids) {
            if (nucleus is Neuroid neuroid) {
                float value = neuroid.outputValue.magnitude;
                if (value > maxValue)
                    maxValue = value;
            }
        }
        float spacing = 400f / nodeCount;
        float margin = 10 + spacing / 2;
        foreach (Nucleus layerNucleus in layer.neuroids) {
            if (layerNucleus is Neuroid layerNeuroid) {
                Vector2Int layerNeuroidPos = this.neuroidPositions[layerNeuroid];
                Vector3 parentPos = new(100 + layerNeuroidPos.x * 100, margin + layerNeuroidPos.y * spacing, 0.1f);

                //int i = 0;
                float inputSpacing = 400f / layerNeuroid.synapses.Count;
                float inputMargin = 10 + inputSpacing / 2;
                // int minStale = 10000;
                //foreach ((Nucleus nucleus, float weight) in layerNeuroid.synapses) {
                foreach (Synapse synapse in layerNeuroid.synapses) {
                    Nucleus nucleus = synapse.nucleus;
                    float weight = synapse.weight;
                    if (this.neuroidPositions.ContainsKey(nucleus)) {
                        Vector2Int inputNeuroidPos = this.neuroidPositions[nucleus];
                        if (inputNeuroidPos.x == layerNeuroidPos.x + 1) {
                            Vector3 pos = new(100 + inputNeuroidPos.x * 100, inputMargin + inputNeuroidPos.y * inputSpacing, 0.0f);

                            float brightness = weight / 10.0f;
                            Handles.color = new Color(brightness, brightness, brightness);
                            Handles.DrawLine(parentPos, pos);
                        }
                    }
                    // if (nucleus is Neuroid neuroid && neuroid.stale < minStale)
                    //     minStale = neuroid.stale;
                }

                // if (layerNeuroid.synapses.Count > 0 && minStale > 2 && layerNeuroid.stale < 3)
                //     Debug.LogWarning($"Strange {minStale} is big duing update");


                float size = 20;
                if (layerNeuroid.isSleeping)
                    Handles.color = Color.darkRed;
                else {
                    float brightness = layerNeuroid.outputValue.magnitude / maxValue;
                    Handles.color = new Color(brightness, brightness, brightness);
                }
                Handles.DrawSolidDisc(parentPos, Vector3.forward, size);
                Vector3 labelPos = parentPos - Vector3.down * (size + 0.2f); // below disc along up axis
                GUIStyle style = new GUIStyle(EditorStyles.label) {
                    alignment = TextAnchor.UpperCenter,
                    normal = { textColor = Color.white },
                    fontStyle = FontStyle.Bold
                };
                Handles.Label(labelPos, layerNeuroid.name, style);

                Rect neuronRect = new(parentPos.x - size, parentPos.y - size, size * 2, size * 2);
                int id = GUIUtility.GetControlID(FocusType.Passive);
                Event e = Event.current;
                EventType et = e.GetTypeForControl(id);
                if (e != null && neuronRect.Contains(e.mousePosition)) {
                    // Process Hover
                    HandleMouseHover(layerNeuroid, neuronRect);
                    // Process click
                    if (et == EventType.MouseDown && e.button == 0) {
                        // Consume the event so the scene doesn't also handle it
                        e.Use();
                        HandleDiscClicked(layerNeuroid);
                    }
                }
                //i++;
            }
        }
    }

    private void HandleMouseHover(Neuroid neuroid, Rect rect) {
        GUIContent tooltip;
        if (neuroid is SensoryNeuroid sensoryNeuroid) {
            tooltip = new(
                $"{sensoryNeuroid.name}" +
                $"\nThing {sensoryNeuroid.receptor.thingType}" +
                $"\nValue: {neuroid.outputValue}");
        }
        else {
            tooltip = new(
                $"{neuroid.name}" +
                $"\nsynapse count {neuroid.synapses.Count}" +
                $"\nValue: {neuroid.outputValue}");
        }

        Vector2 mousePosition = Event.current.mousePosition;

        // Display tooltip with some offset
        Vector2 tooltipSize = GUI.skin.box.CalcSize(tooltip);
        Rect tooltipRect = new Rect(mousePosition.x + 10, mousePosition.y + 10, tooltipSize.x, tooltipSize.y);

        GUI.Box(tooltipRect, tooltip);
    }

    private void HandleDiscClicked(Nucleus nucleus) {
        this.currentNucleus = nucleus;
        BuildLayers();
    }

    protected virtual void OnSceneGUI() {
        NanoBrain brain = target as NanoBrain;
        if (brain == null)
            return;

        Vector3 position = brain.transform.position;
        float radius = 1;


        Handles.DrawWireDisc(position, Vector3.up, radius);    // horizontal circle
        Handles.DrawWireDisc(position, Vector3.right, radius); // X-plane
        Handles.DrawWireDisc(position, Vector3.forward, radius); // Z-plane


        // Debug.DrawRay(brain.transform.position, Vector3.forward, Color.magenta);
        // Handles.color = Color.green;
        // Handles.DrawLine(brain.transform.position, brain.transform.position + Vector3.up);
        Handles.color = Color.yellow;
        Vector3 worldForce = brain.transform.TransformDirection(this.currentNucleus.outputValue.ToVector3());
        //Debug.DrawRay(position, worldForce * 10, Color.yellow);
        Handles.DrawLine(position, position + worldForce * 10);
    }

    #endregion Update
}