using System.Collections.Generic;
using System.Linq;

using UnityEditor;
using UnityEngine;
using UnityEngine.UIElements;

namespace NanoBrain {

    public class ClusterViewer : Editor {

        //public static ClusterViewer previousEditor;
        public static ClusterPrefab previousPrefab;

        public class GraphView : VisualElement {
            protected readonly ClusterPrefab prefab;
            protected SerializedObject serializedBrain;
            protected Nucleus currentNucleus;
            protected Nucleus selectedOutput;

            protected GameObject gameObject;
            private List<NeuroidLayer> layers = new();
            private readonly Dictionary<Nucleus, Vector2Int> neuroidPositions = new();
            private bool expandArray = false;

            protected ClusterPrefab prefabAsset;
            protected VisualElement topMenuContainer;
            protected ScrollView scrollView;
            protected IMGUIContainer graphContainer;
            protected readonly PopupField<string> outputsPopup;

            public ClusterInspector currentEditor;
            //public ClusterViewer previousEditor;

            public enum Mode {
                Focus,
                Full
            }
            public Mode mode = Mode.Focus;

            public GraphView(ClusterPrefab prefab) {
                this.prefab = prefab;

                name = "content";
                style.flexGrow = 1;

                topMenuContainer = new() {
                    style = {
                        flexDirection = FlexDirection.Row,
                        alignItems = Align.Center,
                    }
                };

                EnumField modePopup = new(mode);
                modePopup.style.width = 80;
                modePopup.RegisterValueChangedCallback(OnModeChange);
                topMenuContainer.Add(modePopup);

                List<string> names = this.prefab.outputs.Select(output => output.name).ToList();
                if (names.Count > 0 && names.First() != null) {
                    outputsPopup = new(names, names.First()) {
                        style = { flexGrow = 1 }
                    };
                    outputsPopup.RegisterValueChangedCallback(evt => OnOutputChanged(evt.newValue));
                    topMenuContainer.Add(outputsPopup);
                }

                scrollView = new(ScrollViewMode.Horizontal);
                scrollView.style.position = Position.Absolute;
                scrollView.style.left = 0; scrollView.style.top = 0;
                scrollView.style.right = 0; scrollView.style.bottom = 0;
                //scrollView.style.flexGrow = 1;
                scrollView.horizontalScrollerVisibility = ScrollerVisibility.Auto; // Auto shows when needed
                scrollView.verticalScrollerVisibility = ScrollerVisibility.Hidden;

                graphContainer = new(OnIMGUI);
                //graphContainer.style.position = Position.Relative; // or omit this line
                //graphContainer.style.position = Position.Absolute;
                // graphContainer.style.left = 0; graphContainer.style.top = 0;
                // graphContainer.style.right = 0; graphContainer.style.bottom = 0;
                graphContainer.pickingMode = PickingMode.Position;
                graphContainer.focusable = true;
                //graphContainer.style.width = 1200;
                //graphContainer.style.width = new StyleLength(StyleKeyword.Null); // allow content to determine width

                scrollView.contentContainer.Add(graphContainer);
                Add(scrollView);
                Add(topMenuContainer);


                // Subscribe when added to panel (editor UI ready)
                RegisterCallback<AttachToPanelEvent>(evt => Subscribe());
                RegisterCallback<DetachFromPanelEvent>(evt => Unsubscribe());
            }

            protected virtual void OnModeChange(ChangeEvent<System.Enum> evt) {
                mode = (Mode)evt.newValue;
            }

            protected virtual void OnOutputChanged(string outputName) {
                if (this.currentNucleus.parent != null)
                    // Get nucleus in the parent instance
                    this.selectedOutput = this.currentNucleus.parent.GetNucleus(outputName);
                else
                    // Get nucleus in the prefab
                    this.selectedOutput = this.prefab.GetNucleus(outputName);
                this.currentNucleus = this.selectedOutput;
            }

            bool subscribed = false;
            void Subscribe() {
                if (subscribed) return;
                SceneView.duringSceneGui += OnSceneGUI;
                subscribed = true;
                SceneView.RepaintAll();
            }

            void Unsubscribe() {
                if (!subscribed) return;
                SceneView.duringSceneGui -= OnSceneGUI;
                subscribed = false;
            }

            public void SetGraph(GameObject gameObject, Nucleus nucleus) {
                this.gameObject = gameObject;

                if (Application.isPlaying == false)
                    this.serializedBrain = new SerializedObject(this.prefab);
                this.currentNucleus = nucleus;
                Rebuild(); //inspectorContainer);
                if (outputsPopup != null)
                    OnOutputChanged(outputsPopup.choices[0]);
            }

            void Rebuild() {
                BuildLayers();

                if (this.currentNucleus == null)
                    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");
                }
            }

            protected 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 };

                if (selectedNucleus is Neuron selectedNeuron && selectedNeuron.receivers != null) {
                    foreach (Nucleus receiver in selectedNeuron.receivers) {
                        Nucleus outputNeuroid = receiver;
                        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 };

                if (selectedNucleus.synapses != null) {
                    foreach (Synapse synapse in selectedNucleus.synapses) {
                        Nucleus input = synapse.neuron;
                        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) {
                if (nucleus == null)
                    return;
                layer.neuroids.Add(nucleus);
                // Store its position
                Vector2Int neuroidPosition = new(layer.ix, layer.neuroids.Count - 1);
                neuroidPositions[nucleus] = neuroidPosition;
            }

            public void OnIMGUI() {
                if (currentNucleus == null)
                    return;

                if (Application.isPlaying == false)
                    serializedBrain.Update();

                Handles.BeginGUI();
                DrawGraph();
                Handles.EndGUI();
            }

            #region Graph

            protected virtual void DrawGraph() {
                if (mode == Mode.Focus)
                    DrawFocusGraph();
                else
                    DrawFullGraph();
            }

            #region Full Graph

            protected void DrawFullGraph() {
                //Dag dag = GenerateGraph(this.prefab);
                Dag dag = GenerateGraph(this.selectedOutput);
                Dag.ComputeLayout(dag);
                // Draw edges
                foreach (Dag.Edge e in dag.edges) {
                    Dag.Node from = dag.nodes.FirstOrDefault(x => x.id == e.fromId);
                    Dag.Node to = dag.nodes.FirstOrDefault(x => x.id == e.toId);
                    if (from == null || to == null)
                        continue;

                    Vector2 fromPosition = from.position;
                    Vector2 toPosition = to.position;
                    DrawEdge(fromPosition, toPosition);
                }

                // Draw nodes
                foreach (Dag.Node n in dag.nodes)
                    DrawNucleus(n.nucleus, n.position, 1, n.radius);

                // Determine graph width
                float width = 0;
                float currentNucleusPosition = 0;
                foreach (Dag.Node node in dag.nodes) {
                    if (node.position.x > width)
                        width = node.position.x;
                    if (node.nucleus == currentNucleus)
                        currentNucleusPosition = node.position.x;
                }

                // Resize the graph container to the full graph width
                float margin = 50f;
                graphContainer.style.width = width + 2 * margin;

                // Scroll to the current nucleus
                float viewportWidth = scrollView.layout.width;
                // center currentNucleus in viewport
                float desiredScrollX = currentNucleusPosition - viewportWidth * 0.5f;
                // clamp between 0 and maximum scrollable range
                float maxScrollX = Mathf.Max(0f, graphContainer.resolvedStyle.width - viewportWidth);
                desiredScrollX = Mathf.Clamp(desiredScrollX, 0f, maxScrollX);

                Vector2 current = scrollView.scrollOffset;
                scrollView.scrollOffset = new Vector2(desiredScrollX, current.y);
            }

            public Dag GenerateGraph(Nucleus rootNucleus) {
                Dag dag = new();
                if (rootNucleus == null)
                    return dag;

                int ix = 0;
                Dag.Node receiver = new() {
                    id = ix,
                    //title = nucleus.name,
                    nucleus = rootNucleus
                };
                dag.nodes.Add(receiver);
                ix++;
                DescendGraph(receiver, ref ix, dag);
                return dag;
            }

            private void DescendGraph(Dag.Node receiver, ref int ix, Dag dag) {
                foreach (Synapse synapse in receiver.nucleus.synapses) {
                    Nucleus nucleus = synapse.neuron;
                    if (nucleus.parent != null && nucleus.parent != currentNucleus.parent) {
                        nucleus = nucleus.parent;
                    }
                    string nucleusName = nucleus.name;
                    Dag.Node synapseNode = dag.FindNode(nucleusName);
                    if (synapseNode == null) {
                        synapseNode = new() {
                            id = ix,
                            nucleus = nucleus
                        };
                        dag.nodes.Add(synapseNode);
                    }
                    Dag.Edge edge = new() {
                        fromId = synapseNode.id,
                        toId = receiver.id
                    };
                    dag.edges.Add(edge);
                    ix++;
                    DescendGraph(synapseNode, ref ix, dag);
                }
            }

            #endregion Full Graph

            #region Focus Graph

            protected void DrawFocusGraph() {
                float size = 20;
                Vector3 position = new(150, 210, 0);

                DrawReceivers(this.currentNucleus, position, size);
                DrawSynapses(this.currentNucleus, position, size);

                // Draw selected Nucleus
                if (expandArray) {
                    float maxValue = 1;

                    if (this.currentNucleus is Cluster cluster) {
                        float spacing = 400f / cluster.siblingClusters.Length;
                        float margin = 10 + spacing / 2;
                        float xMin = 150 - size;
                        float xMax = 150 + size;
                        float yMin = 10 + margin - size / 2;
                        float yMax = 400 - margin + size;
                        Vector3[] verts = new Vector3[4] {
                            new(xMin, yMin, 0),
                            new(xMax, yMin, 0),
                            new(xMax, yMax, 0),
                            new(xMin, yMax, 0)
                        };
                        Handles.color = Color.black;
                        Handles.DrawAAConvexPolygon(verts);
                        int row = 0;
                        foreach (Cluster sibling in cluster.siblingClusters) {
                            Vector3 pos = new(150, margin + row * spacing, 0.0f);
                            Handles.color = Color.white;
                            // The selected sibling highlight ring
                            Handles.DrawSolidDisc(pos, Vector3.forward, size + 2);
                            DrawNucleus(sibling, pos, maxValue, size);
                            row++;
                        }
                        GUIStyle style = new(EditorStyles.label) {
                            alignment = TextAnchor.UpperCenter,
                            normal = { textColor = Color.white },
                            fontStyle = FontStyle.Bold,
                        };
                        Vector3 labelPos = new(150, yMax + size + 5, 0);
                        string clusterName = cluster.name;
                        int colonPos = clusterName.IndexOf(":");
                        if (colonPos > 0) {
                            string baseName = clusterName[..colonPos];
                            Handles.Label(labelPos, baseName, style);
                        }
                        else
                            Handles.Label(labelPos, clusterName, style);
                    }
                    else {
                        if (this.currentNucleus is Neuron neuron)
                            maxValue = neuron.outputMagnitude;

                        DrawNucleus(this.currentNucleus, position, maxValue, 20);

                    }
                }
                else {
                    float maxValue = 1;
                    if (this.currentNucleus is Neuron neuron)
                        maxValue = neuron.outputMagnitude;
                    else if (this.currentNucleus is Cluster cluster)
                        maxValue = cluster.defaultOutput.outputMagnitude;
                    DrawNucleus(this.currentNucleus, position, maxValue, 20);
                }
                graphContainer.style.width = 300;
            }

            private void DrawReceivers(Nucleus nucleus, Vector3 parentPos, float size) {
                List<Nucleus> receivers;
                if (nucleus is Neuron neuron)
                    receivers = neuron.receivers;
                else if (nucleus is Cluster cluster)
                    receivers = cluster.CollectReceivers();
                else
                    return;

                int nodeCount = receivers.Count();
                if (nucleus == this.selectedOutput && ClusterViewer.previousPrefab != null) {
                    // Add link to previous editor
                    nodeCount++;
                }

                // Determine the maximum value in this layer
                // This is used to 'scale' the output value colors of the nuclei
                float maxValue = 0;
                foreach (Nucleus receiver in receivers) {
                    if (receiver is Neuron neuroid) {
                        float value = neuroid.outputMagnitude;
                        if (value > maxValue)
                            maxValue = value;
                    }
                }

                // Determine the spacing of the nuclei in the layer
                float spacing = 400f / nodeCount;
                float margin = 10 + spacing / 2;

                int row = 0;
                List<Nucleus[]> drawnArrays = new();
                foreach (Nucleus receiver in receivers) {
                    // if (receiver is Receptor receptor) {
                    //     if (drawnArrays.Contains(receptor.nucleiArray))
                    //         continue;
                    //     drawnArrays.Add(receptor.nucleiArray);
                    // }

                    Nucleus receiverNucleus = receiver;
                    if (receiverNucleus == null)
                        continue;

                    Vector3 pos = new(50, margin + row * spacing, 0.0f);
                    DrawEdge(parentPos, pos);

                    DrawNucleus(receiverNucleus, pos, maxValue, size);
                    row++;
                }
                if (nucleus == this.selectedOutput && ClusterViewer.previousPrefab != null) {
                    Vector3 pos = new(50, margin + row * spacing, 0);
                    DrawEdge(parentPos, pos);
                    DrawClusterPrefab(ClusterViewer.previousPrefab, pos, size);
                }
            }

            private void DrawSynapses(Nucleus nucleus, Vector3 parentPos, float size) {
                int nodeCount = nucleus.synapses.Count;

                // Determine the maximum value in this layer
                // This is used to 'scale' the output value colors of the nuclei
                float maxValue = 0;
                int neuronCount = 0;
                //List<Nucleus[]> drawnArrays = new();
                Cluster[] drawnCluster = null;
                foreach (Synapse synapse in nucleus.synapses) {
                    if (synapse.neuron == null)
                        continue;

                    if (synapse.neuron.parent is Cluster cluster &&
                        //cluster.siblingClusters != null &&
                        synapse.neuron.parent != nucleus.parent) {

                        //if (drawnArrays.Contains(cluster.siblingClusters))
                        if (drawnCluster is not null && cluster.SameSiblingsAs(drawnCluster))
                            continue;
                        //drawnArrays.Add(cluster.siblingClusters);
                        drawnCluster = cluster.siblingClusters;
                    }
                    if (synapse.neuron is Neuron synapseNeuron) {
                        float value = synapseNeuron.outputMagnitude * synapse.weight;
                        // Debug.Log($"{synapse.nucleus.name}: {value}  {length(synapse.nucleus.outputValue)} {synapse.weight}");
                        if (value > maxValue)
                            maxValue = value;
                    }
                    neuronCount++;
                }

                // Determine the spacing of the nuclei in the layer
                float spacing = 400f / neuronCount;
                float margin = 10 + spacing / 2;

                int row = 0;
                //drawnArrays = new();
                drawnCluster = null;
                foreach (Synapse synapse in nucleus.synapses) {
                    if (synapse.neuron is null)
                        continue;

                    if (synapse.neuron.parent is Cluster cluster &&
                        //cluster.siblingClusters != null &&
                        synapse.neuron.parent != nucleus.parent) {

                        // if (drawnArrays.Contains(cluster.siblingClusters))
                        //     continue;
                        // drawnArrays.Add(cluster.siblingClusters);
                        if (drawnCluster is not null && cluster.SameSiblingsAs(drawnCluster))
                            continue;
                        drawnCluster = cluster.siblingClusters;
                    }
                    Vector3 pos = new(250, margin + row * spacing, 0.0f);
                    Handles.color = Color.white;
                    Handles.DrawLine(parentPos, pos);
                    Color color = Color.black;
                    if (Application.isPlaying) {
                        if (maxValue == 0 || !float.IsFinite(maxValue))
                            maxValue = 1;
                        float brightness = 0;
                        if (synapse.neuron is Neuron synapseNeuron)
                            brightness = synapseNeuron.outputMagnitude * synapse.weight / maxValue;
                        color = new Color(brightness, brightness, brightness, 1f);
                    }
                    if (synapse.neuron.parent != null && synapse.neuron.parent != this.currentNucleus.parent) {
                        // the synapse nucleus is part of a subcluster
                        //DrawNucleus(synapse.neuron.parent, pos, maxValue, size, color);
                        DrawNucleus(synapse.neuron, pos, size, color);
                    }
                    else {
                        DrawNucleus(synapse.neuron, pos, size, color);
                    }
                    row++;
                }
            }

            #endregion Focus Graph

            protected void DrawNucleus(Nucleus nucleus, Vector3 position, float maxValue, float size) {
                Color color;
                if (Application.isPlaying) {
                    float brightness = 0;
                    if (nucleus is Neuron neuron)
                        brightness = neuron.outputMagnitude / maxValue;
                    color = new Color(brightness, brightness, brightness, 1f);
                }
                else
                    color = Color.black;
                DrawNucleus(nucleus, position, size, color);
            }

            protected void DrawNucleus(Nucleus nucleus, Vector3 position, float size, Color color) {
                if (nucleus == this.currentNucleus) {
                    // The selected nucleus highlight ring
                    Handles.color = Color.white;
                    Handles.DrawSolidDisc(position, Vector3.forward, size + 2);
                }

                if (nucleus is MemoryCell) {
                    Handles.color = Color.white;
                    Handles.DrawWireDisc(position + Vector3.right * 10, Vector3.forward, size);
                }

                Handles.color = color;
                Handles.DrawSolidDisc(position, Vector3.forward, size);

                Handles.color = Color.white;
                // Position the label in front of the disc
                Vector3 labelPosition = position + (Vector3.forward * 0.1f);

                GUIStyle style = new(EditorStyles.label) {
                    alignment = TextAnchor.MiddleCenter,
                    normal = { textColor = Color.white },
                    fontStyle = FontStyle.Bold,
                };

                if (nucleus.parent is Cluster parentCluster && parentCluster != currentNucleus.parent)
                    DrawCluster(parentCluster, position, color, size);
                else if (nucleus is Cluster cluster)
                    DrawCluster(cluster, position, color, size);

                if (expandArray == false || nucleus != currentNucleus) {
                    // put name below nucleus
                    Vector3 labelPos = position - Vector3.down * (size + 5); // below neuron
                    style.alignment = TextAnchor.UpperCenter;

                    if (nucleus.parent != currentNucleus.parent && nucleus.parent is Cluster parentCluster1) {
                        // This neuron is part of another cluster
                        parentCluster1.name ??= "";
                        string baseName = "";
                        int colonPos = parentCluster1.name.IndexOf(":");
                        if (colonPos > 0 && colonPos < parentCluster1.name.Length - 2)
                            baseName = parentCluster1.name[..colonPos] + ".";
                        else
                            baseName = parentCluster1.name + ".";
                        // if (colonPos > 0 && colonPos < parentCluster1.name.Length - 2) {
                        //     // if it is an array, we should not show the :0 of the first element
                        //     //baseName = baseName[..colonPos];
                        //     Handles.Label(labelPos, baseName + nucleus.name, style);
                        // }
                        // else
                        Handles.Label(labelPos, baseName + nucleus.name, style);
                    }
                    else {
                        nucleus.name ??= "";
                        int colonPos = nucleus.name.IndexOf(":");
                        if (colonPos > 0 && colonPos < nucleus.name.Length - 2) {
                            // if it is an array, we should not show the :0 of the first element
                            string baseName = nucleus.name[..colonPos];
                            Handles.Label(labelPos, baseName, style);
                        }
                        else
                            Handles.Label(labelPos, nucleus.name, style);
                    }
                }

                // Tooltip
                Rect neuronRect = new(position.x - size, position.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(nucleus, neuronRect);
                    // Process click
                    if (e.type == EventType.MouseDown && e.button == 0) {
                        // Consume the event so the scene doesn't also handle it
                        e.Use();
                        if (nucleus is Cluster parentCluster2)
                            HandleClicked(parentCluster2);
                        else
                            HandleClicked(nucleus);
                    }
                }
            }

            private void DrawCluster(Cluster cluster, Vector3 position, Color color, float size) {
                GUIStyle labelTextStyle = new(EditorStyles.label) {
                    normal = { textColor = Color.white },
                    fontStyle = FontStyle.Bold,
                };

                if (expandArray) {
                    // Put array indices above the discs
                    labelTextStyle.alignment = TextAnchor.LowerCenter;
                    Vector3 labelPosition = position + Vector3.down * (size + 5); // below disc 

                    // Strip the instance number in the name
                    int colonPos1 = cluster.name.IndexOf(":");
                    if (colonPos1 > 0) {
                        string extName = cluster.name[(colonPos1 + 2)..];
                        Handles.Label(labelPosition, extName, labelTextStyle);
                    }
                    else
                        Handles.Label(labelPosition, "0", labelTextStyle);
                }
                else {
                    // Put instance count inside the disc
                    labelTextStyle.alignment = TextAnchor.MiddleCenter;
                    Vector3 labelPosition = position + (Vector3.forward * 0.1f);

                    // Adjust text color based on disc color
                    if (color.grayscale > 0.5f)
                        labelTextStyle.normal.textColor = Color.black;
                    else
                        labelTextStyle.normal.textColor = Color.white;

                    if (cluster.instanceCount > 1) {
                        Handles.Label(labelPosition, cluster.instanceCount.ToString(), labelTextStyle);
                        labelTextStyle.normal.textColor = Color.white;
                    }
                    else if (cluster.siblingClusters != null && cluster.siblingClusters.Length > 1) {
                        Handles.Label(labelPosition, cluster.siblingClusters.Length.ToString(), labelTextStyle);
                        labelTextStyle.normal.textColor = Color.white;
                    }
                }

                // Draw a circle around the disc to indicate this is a Cluster
                Handles.color = Color.white;
                Handles.DrawWireDisc(position, Vector3.forward, size + 5);
            }

            protected void DrawClusterPrefab(ClusterPrefab prefab, Vector2 position, float size) {
                Handles.color = Color.black;
                Handles.DrawSolidDisc(position, Vector3.forward, size);
                // Draw a circle around the disc to indicate this is a Cluster
                Handles.color = Color.white;
                Handles.DrawWireDisc(position, Vector3.forward, size + 5);

                // put name below nucleus
                GUIStyle style = new(EditorStyles.label) {
                    alignment = TextAnchor.MiddleCenter,
                    normal = { textColor = Color.white },
                    fontStyle = FontStyle.Bold,
                };
                Vector2 labelPos = position - Vector2.down * (size + 5); // below neuron
                style.alignment = TextAnchor.UpperCenter;
                Handles.Label(labelPos, prefab.name, style);

                Rect neuronRect = new(position.x - size, position.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 click
                    if (e.type == EventType.MouseDown && e.button == 0) {
                        // Consume the event so the scene doesn't also handle it
                        e.Use();
                        Selection.activeObject = prefab;
                        EditorGUIUtility.PingObject(prefab);
                        ClusterViewer.previousPrefab = null;
                        CreateEditor(prefab);
                    }
                }
            }

            protected void DrawEdge(Vector2 from, Vector2 to) {
                Handles.color = Color.white;
                Handles.DrawLine(from, to);
            }

            private void HandleMouseHover(Nucleus nucleus, Rect rect) {
                GUIContent tooltip;
                if (nucleus is Neuron neuron) {
                    tooltip = new(
                        $"{nucleus.name}" +
                        $"\nValue: {neuron.outputMagnitude}");
                }
                else
                    tooltip = new($"{nucleus.name}");

                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);
            }

            protected void HandleClicked(Nucleus nucleus) {
                if (nucleus == this.currentNucleus) {
                    if (Application.isPlaying) {
                        if (nucleus is Cluster)
                            expandArray = !expandArray;
                        else
                            expandArray = false;
                    }
                    else {
                        if (nucleus is Cluster cluster)
                            EditCluster(cluster);
                    }
                }
                else if (nucleus.parent != this.currentNucleus.parent) {
                    // We go to a different cluster
                    // select the cluster, not the neuron in the cluster
                    this.currentNucleus = nucleus.parent;
                    expandArray = false;
                    BuildLayers();
                }
                else {
                    this.currentNucleus = nucleus;
                    expandArray = false;
                    BuildLayers();
                }
            }

            protected void EditCluster(Cluster subCluster) {
                // May be used with storedPrefab...
                Selection.activeObject = subCluster.prefab;
                EditorGUIUtility.PingObject(subCluster.prefab);
                ClusterViewer.previousPrefab = this.prefab;
                ClusterInspector newEditor = CreateEditor(subCluster.prefab) as ClusterInspector;
            }

            #endregion Graph

            void OnSceneGUI(SceneView sceneView) {
                if (this.gameObject != null) {
                    // if (this.currentNucleus is IReceptor receptor) {
                    //     foreach (Nucleus nucleus in receptor.nucleiArray) {
                    //         if (nucleus is Neuron neuron) {
                    //             Vector3 worldVector = this.gameObject.transform.TransformVector(neuron.outputValue);
                    //             Handles.color = Color.yellow;
                    //             Handles.DrawLine(this.gameObject.transform.position, this.gameObject.transform.position + worldVector);
                    //         }
                    //     }
                    // }
                    // else {
                    if (this.currentNucleus is Neuron currentNeuron) {
                        Vector3 worldVector = this.gameObject.transform.TransformVector(currentNeuron.outputValue);
                        Handles.color = Color.yellow;
                        Handles.DrawLine(this.gameObject.transform.position, this.gameObject.transform.position + worldVector);
                    }
                    // }
                }
            }

        }
    }

    public class NeuroidLayer {
        public int ix = 0;
        public List<Nucleus> neuroids = new();
    }

    public class Dag {

        public class Node {
            public int id;
            public Vector2 position;
            public float radius = 20f; // circle radius
            public Nucleus nucleus;
        }

        public class Edge {
            public int fromId;
            public int toId;
        }

        public List<Node> nodes = new();
        public List<Edge> edges = new();

        public Node FindNode(string name, bool justBaseName = true) {
            if (justBaseName) {
                int colonPos = name.IndexOf(":");
                if (colonPos > 0)
                    name = name[..colonPos];
            }
            foreach (Node node in this.nodes) {
                string nodeName = node.nucleus.name;
                if (justBaseName) {
                    int colonPos = nodeName.IndexOf(":");
                    if (colonPos > 0)
                        nodeName = nodeName[..colonPos];
                }
                if (nodeName == name)
                    return node;
            }
            return null;
        }

        public static Node GetNodeById(Dag dag, int id) => dag.nodes.FirstOrDefault(x => x.id == id);

        public static void ComputeLayout(Dag dag) {
            Dictionary<int, List<int>> adjacency = dag.nodes.ToDictionary(n => n.id, n => new List<int>());
            Dictionary<int, int> outdegree = dag.nodes.ToDictionary(node => node.id, n => 0);
            foreach (Edge edge in dag.edges) {
                if (!adjacency.ContainsKey(edge.fromId) || !adjacency.ContainsKey(edge.toId))
                    continue;
                adjacency[edge.fromId].Add(edge.toId);
                outdegree[edge.fromId]++;
            }

            // Kahn's algorithm to compute topological layers (horizontal layers)
            // build parent list (reverse adjacency) and parentIndegree = number of children each parent has
            Dictionary<int, List<int>> parents = dag.nodes.ToDictionary(n => n.id, _ => new List<int>());
            Dictionary<int, int> childCount = dag.nodes.ToDictionary(n => n.id, _ => 0);

            foreach (Edge edge in dag.edges) {
                if (!adjacency.ContainsKey(edge.fromId) || !adjacency.ContainsKey(edge.toId)) continue;
                adjacency[edge.fromId].Add(edge.toId);
                parents[edge.toId].Add(edge.fromId);   // parent of 'to' is 'from'
                childCount[edge.fromId]++;              // outdegree
            }

            Dictionary<int, int> layer = new();
            Queue<int> queue = new(outdegree.Where(kv => kv.Value == 0).Select(kv => kv.Key));
            foreach (int id in queue)
                layer[id] = 0;

            // process parents (reverse traversal)
            while (queue.Count > 0) {
                int u = queue.Dequeue();
                int l = layer[u];
                foreach (int p in parents[u]) {
                    if (!layer.ContainsKey(p) || layer[p] < l + 1)
                        layer[p] = l + 1;
                    childCount[p]--;               // decrement remaining unprocessed children
                    if (childCount[p] == 0)
                        queue.Enqueue(p);
                }
            }

            // Any unreachable nodes -> assign next layers
            int maxLayer = layer.Count > 0 ? layer.Values.Max() : 0;
            foreach (Node node in dag.nodes) {
                if (!layer.ContainsKey(node.id)) {
                    maxLayer++;
                    layer[node.id] = maxLayer;
                }
            }

            // Group nodes by layer (left to right)
            List<List<int>> layers =
                layer.
                GroupBy(kv => kv.Value).
                OrderBy(g => g.Key).
                Select(g => g.Select(x => x.Key).ToList()).
                ToList();

            // Same code without using Linq
            // Build layers dictionary: layerIndex -> List<int> nodeIds
            // Dictionary<int, List<int>> layersDict = new();
            // foreach (KeyValuePair<int, int> kv in layer) {
            //     int nodeId = kv.Key;
            //     int layerIndex = kv.Value;
            //     if (!layersDict.TryGetValue(layerIndex, out List<int> list)) {
            //         list = new List<int>();
            //         layersDict[layerIndex] = list;
            //     }
            //     list.Add(nodeId);
            // }

            // // Determine sorted layer indices
            // List<int> layerIndices = new(layersDict.Keys);
            // layerIndices.Sort(); // ascending order

            // // Build final List<List<int>> in sorted order
            // List<List<int>> layers = new();
            // foreach (int idx in layerIndices) {
            //     layers.Add(layersDict[idx]);
            // }

            float hSpacing = 100f;
            float totalHeight = 400f;

            // Place nodes: x increases with layer index, y spaced within layer
            for (int layerIx = 0; layerIx < layers.Count; layerIx++) {
                List<int> nodeList = layers[layerIx];
                float spacing = totalHeight / nodeList.Count;
                float margin = 10 + spacing / 2;
                for (int i = 0; i < nodeList.Count; i++) {
                    int index = nodeList[i];
                    Node node = GetNodeById(dag, index);
                    if (node == null)
                        continue;
                    float x = hSpacing + layerIx * hSpacing;
                    //float y = 400 - totalHeight / 2f + i * vSpacing;
                    float y = margin + i * spacing;
                    // Debug.Log($"({li}, {i}) -> {x}, {y}");
                    node.position = new Vector2(x, y);
                }
            }

            //Repaint();
        }
    }

}