NanoBrain-unitypackage/Editor/ClusterEditor.cs

628 lines
30 KiB
C#

using System.Collections.Generic;
using System.Linq;
using UnityEditor;
using UnityEngine;
using UnityEngine.UIElements;
namespace NanoBrain {
[CustomEditor(typeof(ClusterPrefab))]
public class ClusterEditor : ClusterViewer {
public override VisualElement CreateInspectorGUI() {
ClusterPrefab prefab = target as ClusterPrefab;
if (prefab != null)
prefab.EnsureInitialization();
serializedObject.Update();
VisualElement root = new();
CreateEditor(root, prefab, null);
serializedObject.ApplyModifiedProperties();
return root;
}
public GraphView CreateEditor(VisualElement root, ClusterPrefab cluster, GameObject gameObject) {
root.style.paddingLeft = 0;
root.style.paddingRight = 0;
root.style.paddingTop = 0;
root.style.paddingBottom = 0;
root.styleSheets.Add(Resources.Load<StyleSheet>("GraphStyles"));
VisualElement mainContainer = new() {
style = {
flexDirection = FlexDirection.Row,
}
};
GraphEditor graphContainer = new(cluster);
graphContainer.style.flexShrink = 0;
graphContainer.style.width = 300;
graphContainer.style.overflow = Overflow.Hidden;
VisualElement inspectorContainer = new() {
name = "inspector",
style = {
minHeight = 450,
width = 300,
flexGrow = 0,
flexDirection = FlexDirection.Row,
}
};
mainContainer.Add(graphContainer);
mainContainer.Add(inspectorContainer);
root.Add(mainContainer);
graphContainer.SetGraph(gameObject, inspectorContainer);
return graphContainer;
}
public class GraphEditor : GraphView {
protected ClusterPrefab prefab;
protected Nucleus currentPrefabNucleus;
protected override Nucleus currentNucleus {
get => base.currentNucleus;
set {
base.currentNucleus = value;
this.currentPrefabNucleus = value != null ? this.prefab.GetNucleus(value.name) : null;
}
}
public GraphEditor(ClusterPrefab prefab) : base(prefab.output.parent) {
this.prefab = prefab;
// In a Prefab editor, no instance exists but we need it for the ClusterViewer.
// So we create a temporary instance
this.currentCluster = new(prefab);
}
public void SetGraph(GameObject gameObject, VisualElement inspectorContainer) {
this.gameObject = gameObject;
if (Application.isPlaying == false)
this.serializedBrain = new SerializedObject(this.prefab);
this.selectedOutput = this.currentCluster.outputs[0];
this.currentNucleus = this.selectedOutput;
//this.currentCluster = this.currentNucleus.parent;
Rebuild(inspectorContainer);
// if (outputsPopup != null)
// OnOutputChanged(outputsPopup.choices[0]);
}
private void Rebuild(VisualElement inspectorContainer) {
if (this.currentNucleus == null) {
inspectorContainer.Clear();
return;
}
string path = AssetDatabase.GetAssetPath(this.prefab); // or known path
this.prefabAsset = AssetDatabase.LoadAssetAtPath<ClusterPrefab>(path);
if (this.prefabAsset == null) {
// create in memory save if it doesn't exist
this.prefabAsset = CreateInstance<ClusterPrefab>();
//Debug.LogError("Cluster Prefab is not found on disk");
}
// DrawInspector(inspectorContainer);
if (inspectorContainer == null)
return;
inspectorContainer.Clear();
if (this.currentNucleus == null)
return;
// create a SerializedObject wrapper so Unity inspector controls work (and Undo)
SerializedObject so = new(prefabAsset);
foreach (Nucleus nucleus in this.prefab.nuclei) {
nucleus.Initialize();
}
this.inspectorIMGUIContainer = new IMGUIContainer(() => InspectorHandler(so));
inspectorContainer.Add(inspectorIMGUIContainer);
}
#region Inspector
private VisualElement inspectorIMGUIContainer;
private bool showSynapses = true;
private bool showActivation = true;
protected bool breakOnWake = false;
protected bool trace = false;
void InspectorHandler(SerializedObject serializedObject) {
bool anythingChanged = false;
if (serializedObject == null || serializedObject.targetObject == null)
return;
serializedObject.Update();
GUIStyle boldTextFieldStyle = new(EditorStyles.textField) {
fontStyle = FontStyle.Bold
};
if (this.currentNucleus == null) {
OutputsInspector(ref anythingChanged);
return;
}
else {
GUIStyle headerStyle = new(EditorStyles.boldLabel) {
alignment = TextAnchor.MiddleLeft,
margin = new RectOffset(10, 0, 4, 4)
};
// Nucleus type
string nucleusType = this.currentNucleus.GetType().Name;
GUILayout.Label(nucleusType, headerStyle);
// Nucleus name
Cluster cluster = this.currentPrefabNucleus as Cluster;
if (cluster != null) {
EditorGUILayout.BeginHorizontal();
if (GUILayout.Button(this.currentNucleus.parent.name))
OnClusterClick(cluster);
EditorGUI.BeginDisabledGroup(true);
EditorGUILayout.TextField(this.currentNucleus.name, boldTextFieldStyle);
EditorGUI.EndDisabledGroup();
if (GUILayout.Button("Reimport"))
ReimportCluster(cluster);
EditorGUILayout.EndHorizontal();
}
else {
string newName = EditorGUILayout.TextField(this.currentNucleus.name, boldTextFieldStyle);
if (newName != this.currentNucleus.name) {
Nucleus prefabNucleus = this.prefab.GetNucleus(this.currentNucleus.name);
prefabNucleus.name = newName;
// This changes it in the temporary cluster instance
this.currentNucleus.name = newName;
this.prefab.RefreshOutputs();
// outputsPopup.choices = this.prefab.outputs.Select(output => output.name).ToList();
anythingChanged = true;
}
}
// Current output value
if (Application.isPlaying) {
if (currentNucleus is Neuron currentNeuron1) {
GUIContent nameLabel = new("Output", currentNeuron1.outputValue.ToString());
EditorGUILayout.FloatField(nameLabel, currentNeuron1.outputMagnitude);
}
else
EditorGUILayout.LabelField(" ");
}
else
EditorGUILayout.LabelField(" ");
// Memory cell
if (this.currentNucleus is MemoryCell memory)
MemoryCellInspector(memory, ref anythingChanged);
// Cluster
else if (cluster != null)
ClusterInspector(cluster, ref anythingChanged);
// Other
else
NucleusInspector(this.currentNucleus, ref anythingChanged);
if (GUILayout.Button("Delete"))
DeleteNucleus(this.currentNucleus);
}
serializedObject.ApplyModifiedProperties();
if (anythingChanged) {
EditorUtility.SetDirty(prefabAsset);
AssetDatabase.SaveAssets();
}
}
protected void OutputsInspector(ref bool anythingChanged) {
GUIStyle headerStyle = new(EditorStyles.boldLabel) {
alignment = TextAnchor.MiddleLeft,
margin = new RectOffset(10, 0, 4, 4)
};
GUILayout.Label("Outputs", headerStyle);
bool connecting = GUILayout.Button("Add Output Neuron");
if (connecting) {
Nucleus newOutput = new Neuron(this.prefab, "New Output");
// Regenerate the temporary clsuter instance
// See also the constructor
this.currentCluster = new(this.prefab);
this.currentNucleus = newOutput;
this.selectedOutput = this.currentNucleus;
}
}
protected void MemoryCellInspector(MemoryCell memoryCell, ref bool anythingChanged) {
memoryCell.staticMemory = EditorGUILayout.Toggle("Static Memory", memoryCell.staticMemory);
NucleusInspector(memoryCell, ref anythingChanged);
}
protected void ClusterInspector(Cluster cluster, ref bool anythingChanged) {
EditorGUILayout.BeginHorizontal();
int instanceCount = cluster.instanceCount;
if (instanceCount <= 1) {
if (cluster.siblingClusters != null && cluster.siblingClusters.Length > 1)
instanceCount = cluster.siblingClusters.Count();
else
instanceCount = 1;
}
EditorGUILayout.IntField("Instances", instanceCount, GUILayout.MinWidth(150));
if (GUILayout.Button("Add")) {
Undo.RecordObject(prefabAsset, "Array add " + prefabAsset.name);
cluster.AddInstance();
anythingChanged = true;
}
if (GUILayout.Button("Del")) {
Undo.RecordObject(prefabAsset, "Array delete " + prefabAsset.name);
cluster.RemoveInstance();
anythingChanged = true;
}
EditorGUILayout.EndHorizontal();
if (GUILayout.Button("Reimport Cluster"))
ReimportCluster(cluster);
}
protected void NucleusInspector(Nucleus nucleus, ref bool anythingChanged) {
SynapsesInspector(ref anythingChanged);
ActivationInspector(ref anythingChanged);
EditorGUILayout.Space();
breakOnWake = EditorGUILayout.Toggle("Break on wake", breakOnWake);
if (breakOnWake && this.currentNucleus is Neuron currentNeuron) {
if (currentNeuron.isSleeping == false)
Debug.Break();
}
trace = EditorGUILayout.Toggle("Trace", trace);
this.currentNucleus.trace = trace;
}
protected void SynapsesInspector(ref bool anythingChanged) {
showSynapses = EditorGUILayout.BeginFoldoutHeaderGroup(showSynapses, "Synapses");
if (showSynapses) {
if (this.currentNucleus is Neuron neuron2) {
Neuron.CombinatorType newCombinator = (Neuron.CombinatorType)EditorGUILayout.EnumPopup("Combinator", neuron2.combinator);
anythingChanged |= newCombinator != neuron2.combinator;
neuron2.combinator = newCombinator;
}
EditorGUIUtility.wideMode = true;
float previousLabelWidth = EditorGUIUtility.labelWidth;
EditorGUIUtility.labelWidth = 100;
Vector3 newBias = EditorGUILayout.Vector3Field("Bias", this.currentNucleus.bias);
if (newBias != this.currentPrefabNucleus.bias) {
anythingChanged |= newBias != this.currentNucleus.bias;
this.currentPrefabNucleus.bias = newBias;
this.currentNucleus.bias = newBias;
}
EditorGUIUtility.labelWidth = previousLabelWidth;
Nucleus[] array = null;
int elementIx = -1;
if (this.currentPrefabNucleus.synapses.Count > 0) {
Synapse[] synapses = this.currentPrefabNucleus.synapses.ToArray();
foreach (Synapse synapse in synapses) {
if (synapse.neuron == null)
continue;
if (array != null) {
if (synapse.neuron.parent is Cluster iCluster && elementIx > 0) {
int thisElementIx = Cluster.GetNucleusIndex(iCluster.clusterNuclei, synapse.neuron);
if (thisElementIx == elementIx)
continue;
else
elementIx = thisElementIx;
}
if (array.Contains(synapse.neuron))
continue;
else if (array.Contains(synapse.neuron.parent))
continue;
}
else {
if (synapse.neuron.parent is Cluster iReceptor) {
array = iReceptor.siblingClusters;
if (iReceptor is Cluster iCluster)
elementIx = Cluster.GetNucleusIndex(iCluster.clusterNuclei, synapse.neuron);
}
}
EditorGUILayout.Space();
if (Application.isPlaying) {
if (synapse.neuron is Neuron synapseNeuron) {
Vector3 value = synapseNeuron.outputValue * synapse.weight;
GUIContent synapseValueLabel = new(synapse.neuron.name, synapseNeuron.outputValue.ToString());
EditorGUILayout.FloatField(synapseValueLabel, synapseNeuron.outputMagnitude);
}
}
else {
EditorGUILayout.BeginHorizontal();
if (synapse.neuron.clusterPrefab != this.currentNucleus.clusterPrefab) {
// If it is a different cluster
GUIStyle labelStyle = new(GUI.skin.label);
float labelWidth = 200;
if (synapse.neuron.clusterPrefab != null) {
labelWidth = labelStyle.CalcSize(new GUIContent($"{synapse.neuron.clusterPrefab.name}.")).x;
GUILayout.Label($"{synapse.neuron.clusterPrefab.name}", GUILayout.Width(labelWidth));
}
string[] options = synapse.neuron.clusterPrefab.nuclei.Select(n => n.name).ToArray();
int selectedIndex = System.Array.IndexOf(options, synapse.neuron.name);
int newIndex = EditorGUILayout.Popup(selectedIndex, options);
if (newIndex != selectedIndex) {
// Nucleus selectedNucleus = synapse.neuron.parent.clusterNuclei[newIndex];
// Neuron newNeuron = selectedNucleus as Neuron;
Neuron newNeuron = synapse.neuron.clusterPrefab.nuclei[newIndex] as Neuron;
ChangeSynapse(synapse, newNeuron);
}
}
else
GUILayout.Label(synapse.neuron.name);
bool disconnecting = GUILayout.Button("Disconnect", GUILayout.Width(80));
if (disconnecting && synapse.neuron is Neuron synapseNeuron) {
synapseNeuron.RemoveReceiver(this.currentNucleus);
this.prefab.GarbageCollection();
anythingChanged = true;
}
EditorGUILayout.EndHorizontal();
}
EditorGUI.indentLevel++;
float newWeight = EditorGUILayout.FloatField("Weight", synapse.weight);
if (newWeight != synapse.weight) {
synapse.weight = newWeight;
anythingChanged = true;
}
EditorGUI.indentLevel--;
}
}
EditorGUILayout.Space();
anythingChanged |= ConnectNucleus(this.prefab, this.currentNucleus);
anythingChanged |= AddSynapse(this.prefab, this.currentNucleus);
}
EditorGUILayout.EndFoldoutHeaderGroup();
}
protected void ActivationInspector(ref bool anythingChanged) {
EditorGUILayout.Space();
showActivation = EditorGUILayout.BeginFoldoutHeaderGroup(showActivation, "Activation");
if (showActivation) {
if (this.currentNucleus is Neuron neuron) {
if (this.currentNucleus is not MemoryCell) {
EditorGUILayout.BeginHorizontal();
EditorGUILayout.LabelField("Activation Curve", GUILayout.MinWidth(60));
if (neuron.curveMax > 0)
EditorGUILayout.CurveField(neuron.curve, Color.cyan, new Rect(0, 0, 1, neuron.curveMax), GUILayout.Width(40));
else
EditorGUILayout.CurveField(neuron.curve, Color.cyan, new Rect(0, neuron.curveMax, 1, -neuron.curveMax), GUILayout.Width(40));
Neuron.ActivationType newPreset = (Neuron.ActivationType)EditorGUILayout.EnumPopup(neuron.curvePreset, GUILayout.MinWidth(50));
anythingChanged |= newPreset != neuron.curvePreset;
neuron.curvePreset = newPreset;
EditorGUILayout.EndHorizontal();
}
// if (neuron is Receptor receptor2) {
// if (receptor2.nucleiArray == null || receptor2.nucleiArray.Count() == 0)
// receptor2.array = new NucleusArray(neuron);
// }
}
EditorGUILayout.Space();
}
EditorGUILayout.EndFoldoutHeaderGroup();
}
#region Synapses
protected virtual void AddInput(Nucleus.Type selectedType, Nucleus nucleus) {
switch (selectedType) {
case Nucleus.Type.Neuron:
AddNeuronInput(nucleus);
break;
case Nucleus.Type.MemoryCell:
AddMemoryCellInput(nucleus);
break;
case Nucleus.Type.Cluster:
AddClusterInput(nucleus);
break;
default:
break;
}
}
protected virtual void AddNeuronInput(Nucleus nucleus) {
Neuron newNeuroid = new(this.prefab, "New neuron");
newNeuroid.AddReceiver(nucleus);
this.currentNucleus = newNeuroid;
}
protected virtual void AddMemoryCellInput(Nucleus nucleus) {
MemoryCell newMemory = new(this.prefab, "New memory cell");
newMemory.AddReceiver(nucleus);
this.currentNucleus = newMemory;
}
protected virtual void AddClusterInput(Nucleus nucleus) {
ClusterPickerWindow.ShowPicker(brain => OnClusterPicked(nucleus, brain), "Select Cluster");
}
private void OnClusterPicked(Nucleus nucleus, ClusterPrefab selectedPrefab) {
Cluster subclusterInstance = new(selectedPrefab, this.prefab);
subclusterInstance.defaultOutput.AddReceiver(nucleus);
}
private void ReimportCluster(Cluster subCluster) {
if (subCluster.siblingClusters == null || subCluster.siblingClusters.Length <= 0) {
Cluster reimportedCluster = new(subCluster.prefab, this.prefab);
subCluster.MoveReceivers(reimportedCluster);
// subcluster should be garbage now...
this.currentNucleus = reimportedCluster;
}
else {
this.currentNucleus = null;
List<Cluster> newSiblingsList = new();
foreach (Cluster sibling in subCluster.siblingClusters) {
Cluster reimportedCluster = new(sibling.prefab, this.prefab) {
name = sibling.name
};
sibling.MoveReceivers(reimportedCluster);
newSiblingsList.Add(reimportedCluster);
// make the first reimportedCluster the new current nucleus
this.currentNucleus ??= reimportedCluster;
}
Cluster[] newSiblings = newSiblingsList.ToArray();
foreach (Cluster sibling in newSiblings)
sibling.siblingClusters = newSiblings;
}
}
int selectedConnectNucleus = -1;
// Connect to another nucleus
protected virtual bool ConnectNucleus(ClusterPrefab cluster, Nucleus nucleusToConnect) {
if (cluster == null)
return false;
IEnumerable<Nucleus> synapseNuclei = this.currentNucleus.synapses
.Where(synapse => synapse.neuron != null)
.Select(synapse => synapse.neuron);
IEnumerable<Nucleus> nuclei = cluster.nuclei
.Except(synapseNuclei);
IEnumerable<string> nucleiNames = nuclei
.Select(n => {
int idx = n.name.IndexOf(':');
return idx < 0 ? n.name : n.name[..idx];
})
.Distinct();
string[] names = nucleiNames.ToArray();
EditorGUILayout.BeginHorizontal();
selectedConnectNucleus = EditorGUILayout.Popup(selectedConnectNucleus, names);
bool connecting = GUILayout.Button("Connect", GUILayout.Width(80));
EditorGUILayout.EndHorizontal();
if (connecting) {
Nucleus nucleus = nuclei.ElementAt(selectedConnectNucleus);
if (nucleus is Cluster subCluster)
subCluster.AddArrayReceiver(this.currentNucleus);
else if (nucleus is Neuron neuron)
neuron.AddReceiver(this.currentNucleus);
}
return connecting;
}
protected virtual void DeleteNucleus(Nucleus nucleus) {
if (nucleus == null)
return;
if (nucleus is Neuron neuron) {
foreach (Nucleus receiver in neuron.receivers) {
if (receiver != null) {
this.currentNucleus = receiver;
break;
}
}
}
this.currentCluster.DeleteNucleus(nucleus);//clusterNuclei.Remove(nucleus);
// this.prefab.nuclei.Remove(nucleus);
// Neuron.Delete(nucleus);
this.prefab.RefreshOutputs();
this.currentNucleus = this.prefab.output;
this.selectedOutput = this.currentNucleus;
}
Nucleus.Type selectedType = Nucleus.Type.None;
protected virtual bool AddSynapse(ClusterPrefab cluster, Nucleus nucleus) {
if (cluster == null)
return false;
EditorGUILayout.BeginHorizontal();
selectedType = (Nucleus.Type)EditorGUILayout.EnumPopup(selectedType);
bool connecting = GUILayout.Button("Add", GUILayout.Width(80));
EditorGUILayout.EndHorizontal();
if (connecting) {
AddInput(selectedType, this.currentNucleus);
}
return connecting;
// if (selectedType == Nucleus.Type.None)
// return false;
// AddInput(selectedType, this.currentNucleus);
// return true;
}
protected virtual void ChangeSynapse(Synapse synapse, Neuron newNucleus) {
Neuron synapseNeuron = synapse.neuron;
if (synapse.neuron.parent is Cluster subCluster && subCluster.prefab != this.prefab) {
// if (synapse.neuron.parent is ClusterReceptor receptor) {
// // the new nucleus is part of a (cluster) receptor,
// // so we have to change all synapses to this nucleus array elements
// int oldNucleusIx = Cluster.GetNucleusIndex(subCluster.clusterNuclei, synapse.neuron);
// int newNucleusIx = Cluster.GetNucleusIndex(subCluster.clusterNuclei, newNucleus);
// foreach (Nucleus element in receptor.nucleiArray) {
// if (element is not ClusterReceptor clusterReceptor)
// continue;
// // Get the same neuron as the synapse.nucleus in a different element
// // of the ClusterReceptor array
// Nucleus oldElementNucleus = clusterReceptor.clusterNuclei[oldNucleusIx];
// if (oldElementNucleus is not Neuron oldElementNeuron)
// continue;
// // Get the same neuron as newNucleus in a different element
// // of the ClusterReceptor array
// Nucleus newElementNucleus = clusterReceptor.clusterNuclei[newNucleusIx];
// if (newElementNucleus is not Neuron newElementNeuron)
// continue;
// oldElementNeuron.RemoveReceiver(this.currentNucleus);
// newElementNeuron.AddReceiver(this.currentNucleus);
// // Now find the synapse which pointed to the old Neuron
// // Synapse synapseForUpdate = this.currentNucleus.GetSynapse(oldElementNeuron);
// // synapseForUpdate.nucleus = newElementNeuron;
// }
// }
// else {
// it is a neuron in a subcluster
synapseNeuron.RemoveReceiver(this.currentPrefabNucleus);
newNucleus.AddReceiver(this.currentPrefabNucleus);
// }
}
else {
synapseNeuron.RemoveReceiver(this.currentNucleus);
newNucleus.AddReceiver(this.currentNucleus);
}
}
protected virtual void DisconnectNucleus(Neuron nucleus) {
if (this.currentNucleus.clusterPrefab == null)
return;
string[] names = this.currentNucleus.synapses.Select(synapse => synapse.neuron.name).ToArray();
int selectedIndex = -1;
selectedIndex = EditorGUILayout.Popup("Disconnect from", selectedIndex, names);
if (selectedIndex >= 0 && selectedIndex < this.currentNucleus.clusterPrefab.nuclei.Count) {
Synapse synapse = this.currentNucleus.synapses[selectedIndex];
Neuron synapseNeuron = synapse.neuron as Neuron;
synapseNeuron.RemoveReceiver(this.currentNucleus);
}
}
#endregion Synapses
#endregion Inspector
}
}
}