NanoBrain-unitypackage/Editor/DAGWindow.cs

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

// Simple DAG data model
// [System.Serializable]
// public class DagNode
// {
//     public int id;
//     public string title;
//     public Vector2 position;
//     public float radius = 36f; // circle radius
// }

// [System.Serializable]
// public class DagEdge
// {
//     public int fromId;
//     public int toId;
// }

public class DAGEditorWindow : EditorWindow
{
    List<DagNode> nodes = new List<DagNode>();
    List<DagEdge> edges = new List<DagEdge>();

    Vector2 pan = Vector2.zero;
    float zoom = 1.0f;
    const float minZoom = 0.5f;
    const float maxZoom = 2.0f;

    GUIStyle labelStyle;
    int selectedNodeId = -1;

    Vector2 dragStart;
    bool draggingNode = false;
    int draggingNodeId = -1;

    [MenuItem("Window/DAG Viewer (LR, Circles)")]
    public static void ShowWindow()
    {
        var w = GetWindow<DAGEditorWindow>("DAG Viewer (LR)");
        w.minSize = new Vector2(500, 300);
    }

    void OnEnable()
    {
        labelStyle = new GUIStyle(EditorStyles.label);
        labelStyle.alignment = TextAnchor.MiddleCenter;
        labelStyle.normal.textColor = Color.white;
        labelStyle.fontStyle = FontStyle.Bold;

        if (nodes.Count == 0)
            CreateSampleGraph();

        ComputeLeftToRightLayout();
    }

    void CreateSampleGraph()
    {
        nodes.Clear();
        edges.Clear();

        nodes.Add(new DagNode() { id = 0, title = "In1" });
        nodes.Add(new DagNode() { id = 1, title = "In2" });
        nodes.Add(new DagNode() { id = 2, title = "A" });
        nodes.Add(new DagNode() { id = 3, title = "B" });
        nodes.Add(new DagNode() { id = 4, title = "C" });
        nodes.Add(new DagNode() { id = 5, title = "Out1" });
        nodes.Add(new DagNode() { id = 6, title = "Out2" });

        edges.Add(new DagEdge() { fromId = 0, toId = 2 });
        edges.Add(new DagEdge() { fromId = 1, toId = 2 });
        edges.Add(new DagEdge() { fromId = 2, toId = 3 });
        edges.Add(new DagEdge() { fromId = 2, toId = 4 });
        edges.Add(new DagEdge() { fromId = 3, toId = 5 });
        edges.Add(new DagEdge() { fromId = 4, toId = 6 });
    }

    void OnGUI()
    {
        HandleInput();

        Rect rect = new Rect(0, 0, position.width, position.height);
        EditorGUI.DrawRect(rect, new Color(0.11f, 0.11f, 0.11f));

        Matrix4x4 oldMatrix = GUI.matrix;
        Vector2 origin = new Vector2(position.width / 2, position.height / 2);
        GUI.matrix = Matrix4x4.TRS(origin + pan, Quaternion.identity, Vector3.one * zoom) *
                     Matrix4x4.TRS(-origin, Quaternion.identity, Vector3.one);

        // Draw edges first
        foreach (var e in edges)
        {
            var from = GetNodeById(e.fromId);
            var to = GetNodeById(e.toId);
            if (from == null || to == null) continue;
            DrawEdgeCircleNodes(from, to);
        }

        // Draw nodes (circles)
        foreach (var n in nodes)
        {
            DrawNodeCircle(n);
        }

        GUI.matrix = oldMatrix;

        // Footer toolbar
        GUILayout.FlexibleSpace();
        EditorGUILayout.BeginHorizontal(EditorStyles.toolbar);
        if (GUILayout.Button("Fit", EditorStyles.toolbarButton)) FitToView();
        if (GUILayout.Button("Layout LR", EditorStyles.toolbarButton)) ComputeLeftToRightLayout();
        if (GUILayout.Button("Add Node", EditorStyles.toolbarButton))
        {
            AddNode("N" + nodes.Count);
            ComputeLeftToRightLayout();
        }
        if (GUILayout.Button("Add Edge (selected->new)", EditorStyles.toolbarButton))
        {
            if (selectedNodeId != -1)
            {
                var newNode = AddNode("N" + nodes.Count);
                edges.Add(new DagEdge() { fromId = selectedNodeId, toId = newNode.id });
                ComputeLeftToRightLayout();
            }
        }
        EditorGUILayout.EndHorizontal();
    }

    void HandleInput()
    {
        Event e = Event.current;

        // Zoom with scroll
        if (e.type == EventType.ScrollWheel)
        {
            float oldZoom = zoom;
            float delta = -e.delta.y * 0.01f;
            zoom = Mathf.Clamp(zoom + delta, minZoom, maxZoom);
            Vector2 mouse = e.mousePosition;
            pan += (mouse - new Vector2(position.width / 2, position.height / 2)) * (1 - zoom / oldZoom);
            e.Use();
        }

        // Pan with middle or right+ctrl drag
        if (e.type == EventType.MouseDrag && (e.button == 2 || (e.button == 1 && e.control)))
        {
            pan += e.delta;
            e.Use();
        }

        // Node dragging & selection (convert mouse to graph space)
        Vector2 graphMouse = ScreenToGraph(e.mousePosition);
        if (e.type == EventType.MouseDown && e.button == 0)
        {
            int hit = HitTestNode(graphMouse);
            if (hit != -1)
            {
                selectedNodeId = hit;
                draggingNode = true;
                draggingNodeId = hit;
                dragStart = graphMouse;
                e.Use();
            }
            else
            {
                selectedNodeId = -1;
            }
        }

        if (draggingNode && draggingNodeId != -1)
        {
            if (e.type == EventType.MouseDrag && e.button == 0)
            {
                Vector2 graphDelta = e.delta / zoom;
                var n = GetNodeById(draggingNodeId);
                if (n != null)
                {
                    n.position += graphDelta;
                    Repaint();
                    e.Use();
                }
            }
            if (e.type == EventType.MouseUp && e.button == 0)
            {
                draggingNode = false;
                draggingNodeId = -1;
                e.Use();
            }
        }
    }

    DagNode AddNode(string title)
    {
        int nextId = nodes.Count > 0 ? nodes.Max(n => n.id) + 1 : 0;
        var n = new DagNode() { id = nextId, title = title, position = Vector2.zero };
        nodes.Add(n);
        return n;
    }

    DagNode GetNodeById(int id) => nodes.FirstOrDefault(x => x.id == id);

    void DrawNodeCircle(DagNode n)
    {
        Vector2 center = n.position;
        float r = n.radius;
        Rect nodeRect = new Rect(center.x - r, center.y - r, r * 2, r * 2);

        // circle background
        Color bg = (n.id == selectedNodeId) ? new Color(0.15f, 0.5f, 0.9f) : new Color(0.2f, 0.2f, 0.2f);
        EditorGUI.DrawRect(nodeRect, bg);

        // anti-aliased circle outline
        Handles.color = Color.white * 0.9f;
        Handles.DrawAAPolyLine(3f / zoom, GetCircleOutlinePoints(center, r, 48).ToArray());

        // label
        Vector2 labelPos = center - new Vector2(0, 8);
        GUI.Label(new Rect(labelPos.x - r, labelPos.y - 8, r * 2, 18), n.title, labelStyle);
    }

    List<Vector3> GetCircleOutlinePoints(Vector2 center, float radius, int segments)
    {
        var pts = new List<Vector3>(segments + 1);
        for (int i = 0; i <= segments; i++)
        {
            float a = (float)i / segments * Mathf.PI * 2f;
            pts.Add(new Vector3(center.x + Mathf.Cos(a) * radius, center.y + Mathf.Sin(a) * radius, 0));
        }
        return pts;
    }

    void DrawEdgeCircleNodes(DagNode from, DagNode to)
    {
        Vector2 a = from.position;
        Vector2 b = to.position;
        if (a == b) return;

        // Compute edge line that starts/ends at circle circumferences
        Vector2 dir = (b - a).normalized;
        Vector2 start = a + dir * from.radius;
        Vector2 end = b - dir * to.radius;

        // Use a simple curved line: start -> control -> end (bezier)
        Vector2 control = new Vector2((start.x + end.x) / 2f, (start.y + end.y) / 2f);
        // Slight vertical offset to separate overlapping lines based on node ids
        float offset = ((from.id * 7 + to.id * 11) % 7 - 3) * 6f / zoom;
        control += new Vector2(0, offset);

        Handles.color = Color.white * 0.9f;
        Handles.DrawAAPolyLine(3f / zoom, 20, GetBezierPoints(start, control, end, 24).ToArray());

        // Arrow at end pointing towards 'b'
        DrawArrowHead(end - dir * 2f, end, 10f / zoom, 12f / zoom, Color.white);
    }

    List<Vector3> GetBezierPoints(Vector2 p0, Vector2 p1, Vector2 p2, int seg)
    {
        var pts = new List<Vector3>(seg + 1);
        for (int i = 0; i <= seg; i++)
        {
            float t = (float)i / seg;
            Vector2 p = (1 - t) * (1 - t) * p0 + 2 * (1 - t) * t * p1 + t * t * p2;
            pts.Add(new Vector3(p.x, p.y, 0));
        }
        return pts;
    }

    void DrawArrowHead(Vector2 from, Vector2 to, float headWidth, float headLength, Color color)
    {
        Vector2 dir = (to - from).normalized;
        if (dir == Vector2.zero) return;
        Vector2 right = new Vector2(-dir.y, dir.x);

        Vector3 p1 = to;
        Vector3 p2 = to - dir * headLength + right * headWidth * 0.5f;
        Vector3 p3 = to - dir * headLength - right * headWidth * 0.5f;

        Handles.color = color;
        Handles.DrawAAConvexPolygon(p1, p2, p3);
    }

    // Left-to-right layered layout (sources on the left, sinks on the right)
    void ComputeLeftToRightLayout()
    {
        // build adjacency and indegree
        var adj = nodes.ToDictionary(n => n.id, n => new List<int>());
        var indeg = nodes.ToDictionary(n => n.id, n => 0);
        foreach (var e in edges)
        {
            if (!adj.ContainsKey(e.fromId) || !adj.ContainsKey(e.toId)) continue;
            adj[e.fromId].Add(e.toId);
            indeg[e.toId]++;
        }

        // Kahn's algorithm to compute topological layers (horizontal layers)
        Dictionary<int, int> layer = new Dictionary<int, int>();
        Queue<int> q = new Queue<int>(indeg.Where(kv => kv.Value == 0).Select(kv => kv.Key));
        foreach (var id in q) layer[id] = 0;

        while (q.Count > 0)
        {
            int u = q.Dequeue();
            int l = layer[u];
            foreach (var v in adj[u])
            {
                // prefer placing v at least one layer after u
                if (!layer.ContainsKey(v) || layer[v] < l + 1) layer[v] = l + 1;
                indeg[v]--;
                if (indeg[v] == 0) q.Enqueue(v);
            }
        }

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

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

        // Layout parameters (horizontal spacing drives left->right)
        float hSpacing = 220f;
        float vSpacing = 120f;

        // Place nodes: x increases with layer index, y spaced within layer
        for (int li = 0; li < layers.Count; li++)
        {
            var lst = layers[li];
            float totalHeight = (lst.Count - 1) * vSpacing;
            for (int i = 0; i < lst.Count; i++)
            {
                int id = lst[i];
                var n = GetNodeById(id);
                if (n == null) continue;
                float x = li * hSpacing;
                float y = -totalHeight / 2f + i * vSpacing;
                n.position = new Vector2(x, y);
            }
        }

        Repaint();
    }

    void FitToView()
    {
        if (nodes.Count == 0) return;
        Rect bounds = new Rect(nodes[0].position - Vector2.one * nodes[0].radius, Vector2.one * nodes[0].radius * 2f);
        foreach (var n in nodes)
            bounds = RectUnion(bounds, new Rect(n.position - Vector2.one * n.radius, Vector2.one * n.radius * 2f));

        Vector2 center = bounds.center;
        pan = -center;
        zoom = 1.0f;
        Repaint();
    }

    static Rect RectUnion(Rect a, Rect b)
    {
        float xMin = Mathf.Min(a.xMin, b.xMin);
        float xMax = Mathf.Max(a.xMax, b.xMax);
        float yMin = Mathf.Min(a.yMin, b.yMin);
        float yMax = Mathf.Max(a.yMax, b.yMax);
        return Rect.MinMaxRect(xMin, yMin, xMax, yMax);
    }

    Vector2 ScreenToGraph(Vector2 screenPos)
    {
        Vector2 origin = new Vector2(position.width / 2, position.height / 2);
        // invert the GUI.matrix transform (approx for current simple transforms)
        return (screenPos - (origin + pan)) / zoom + origin * (1 - 1 / zoom);
    }

    int HitTestNode(Vector2 graphPos)
    {
        // returns node id under point or -1
        for (int i = nodes.Count - 1; i >= 0; i--)
        {
            var n = nodes[i];
            if ((graphPos - n.position).sqrMagnitude <= n.radius * n.radius) return n.id;
        }
        return -1;
    }
}