DataModel.cpp

#include "DataModel.h"
#include <cstdio>

namespace GraphData {

Node MakeGridNode(int id) {
    Node node;
    node.id = id;
    node.name = "Node " + std::to_string(id);

    // Ports (matching current setup: 5 inputs, 5 outputs)
    node.inputs.push_back(Port("Float", PortType::FLOAT));
    node.inputs.push_back(Port("Int", PortType::INT));
    node.inputs.push_back(Port("Float", PortType::FLOAT));
    node.inputs.push_back(Port("Float", PortType::FLOAT));
    node.inputs.push_back(Port("Int", PortType::INT));

    node.outputs.push_back(Port("Float", PortType::FLOAT));
    node.outputs.push_back(Port("Int", PortType::INT));
    node.outputs.push_back(Port("Float", PortType::FLOAT));
    node.outputs.push_back(Port("Float", PortType::FLOAT));
    node.outputs.push_back(Port("Int", PortType::INT));

    return node;
}

void DataModel::CreateNodesGrid(int nodeCount) {
    Clear();

    // Generate nodes
    for (int i = 0; i < nodeCount; i++) {
        Node node = MakeGridNode(nextNodeId_++);

        // Add to model
        size_t index = nodes_.size();
        nodes_.push_back(node);
        nodeIdToIndex_[node.id] = index;

        // Connect to previous node with 5 edges (matching current setup)
        if (i > 0) {
            int prevNodeId = nodes_[i - 1].id;

            for (int edge = 0; edge < 5; edge++) {
                Connection conn(prevNodeId, edge, node.id, edge);
                connections_.push_back(conn);
            }
        }
    }

    printf("DataModel: Created %d nodes with %zu connections\n",
           nodeCount, connections_.size());
}

int DataModel::AddNode(const Node& node) {
    Node newNode = node;
    if (newNode.id == -1) {
        newNode.id = nextNodeId_++;
    }

    size_t index = nodes_.size();
    nodes_.push_back(newNode);
    nodeIdToIndex_[newNode.id] = index;

    return newNode.id;
}

const Node* DataModel::GetNode(int nodeId) const {
    auto it = nodeIdToIndex_.find(nodeId);
    if (it == nodeIdToIndex_.end()) {
        return nullptr;
    }
    return &nodes_[it->second];
}

Node* DataModel::GetNodeMutable(int nodeId) {
    auto it = nodeIdToIndex_.find(nodeId);
    if (it == nodeIdToIndex_.end()) {
        return nullptr;
    }
    return &nodes_[it->second];
}

void DataModel::AddConnection(const Connection& conn) {
    if (!IsValidConnection(conn)) {
        printf("Warning: Invalid connection attempted\n");
        return;
    }
    connections_.push_back(conn);
}

bool DataModel::IsValidConnection(const Connection& conn) const {
    const Node* srcNode = GetNode(conn.sourceNodeId);
    const Node* dstNode = GetNode(conn.targetNodeId);

    if (!srcNode || !dstNode) {
        return false;
    }

    if (conn.sourcePortIndex < 0 || conn.sourcePortIndex >= (int)srcNode->outputs.size()) {
        return false;
    }

    if (conn.targetPortIndex < 0 || conn.targetPortIndex >= (int)dstNode->inputs.size()) {
        return false;
    }

    // Could add type checking here:
    // const Port& srcPort = srcNode->outputs[conn.sourcePortIndex];
    // const Port& dstPort = dstNode->inputs[conn.targetPortIndex];
    // return srcPort.type == dstPort.type;

    return true;
}


void DataModel::CreateTestHierarchy() {
    Clear();

    // AI Model Profiler-style graph with neural network operations

    // Regular Node 0
    Node node0;
    node0.id = nextNodeId_++;
    node0.name = "Input [224x224x3]";
    node0.isExpandable = false;
    node0.parentId = -1;
    node0.outputs.push_back(Port("float32", PortType::FLOAT));

    // Expandable Node 1 (Convolution Block - 4 children with multiple ports)
    Node node1;
    node1.id = nextNodeId_++;
    node1.name = "Conv2D Block";
    node1.isExpandable = true;
    node1.parentId = -1;
    node1.inputs.push_back(Port("float32", PortType::FLOAT));
    node1.outputs.push_back(Port("float32", PortType::FLOAT));

    // Children of Node 1 (Complex convolution internals)
    Node node1_child0;
    node1_child0.id = nextNodeId_++;
    node1_child0.name = "Conv2D [3x3]";
    node1_child0.isExpandable = false;
    node1_child0.parentId = node1.id;
    node1_child0.inputs.push_back(Port("input float32", PortType::FLOAT));
    node1_child0.inputs.push_back(Port("weights float32", PortType::FLOAT));
    node1_child0.inputs.push_back(Port("bias float32", PortType::FLOAT));
    node1_child0.outputs.push_back(Port("float32", PortType::FLOAT));
    node1_child0.outputs.push_back(Port("stats float64", PortType::FLOAT));
    node1_child0.parentInputPortMap = { 0, -1, -1 };  // First input maps to parent
    node1_child0.parentOutputPortMap = { -1, -1 }; // Both internal

    Node node1_child1;
    node1_child1.id = nextNodeId_++;
    node1_child1.name = "BatchNorm";
    node1_child1.isExpandable = false;
    node1_child1.parentId = node1.id;
    node1_child1.inputs.push_back(Port("float32", PortType::FLOAT));
    node1_child1.inputs.push_back(Port("stats float64", PortType::FLOAT));
    node1_child1.outputs.push_back(Port("float32", PortType::FLOAT));
    node1_child1.outputs.push_back(Port("mean float64", PortType::FLOAT));
    node1_child1.outputs.push_back(Port("var float64", PortType::FLOAT));
    node1_child1.parentInputPortMap = { -1, -1 }; // Internal
    node1_child1.parentOutputPortMap = { -1, -1, -1 }; // Internal

    Node node1_child2;
    node1_child2.id = nextNodeId_++;
    node1_child2.name = "ReLU";
    node1_child2.isExpandable = false;
    node1_child2.parentId = node1.id;
    node1_child2.inputs.push_back(Port("float32", PortType::FLOAT));
    node1_child2.outputs.push_back(Port("float32", PortType::FLOAT));
    node1_child2.parentInputPortMap = { -1 }; // Internal
    node1_child2.parentOutputPortMap = { -1 }; // Internal

    Node node1_child3;
    node1_child3.id = nextNodeId_++;
    node1_child3.name = "Dropout [0.2]";
    node1_child3.isExpandable = false;
    node1_child3.parentId = node1.id;
    node1_child3.inputs.push_back(Port("float32", PortType::FLOAT));
    node1_child3.outputs.push_back(Port("float32", PortType::FLOAT));
    node1_child3.parentInputPortMap = { -1 }; // Internal
    node1_child3.parentOutputPortMap = { 0 }; // Maps to parent output

    node1.childNodeIds.push_back(node1_child0.id);
    node1.childNodeIds.push_back(node1_child1.id);
    node1.childNodeIds.push_back(node1_child2.id);
    node1.childNodeIds.push_back(node1_child3.id);

    // Regular Node 2
    Node node2;
    node2.id = nextNodeId_++;
    node2.name = "MaxPool2D [2x2]";
    node2.isExpandable = false;
    node2.parentId = -1;
    node2.inputs.push_back(Port("float32", PortType::FLOAT));
    node2.outputs.push_back(Port("float32", PortType::FLOAT));

    // Regular Node 3
    Node node3;
    node3.id = nextNodeId_++;
    node3.name = "Conv2D [5x5]";
    node3.isExpandable = false;
    node3.parentId = -1;
    node3.inputs.push_back(Port("float32", PortType::FLOAT));
    node3.outputs.push_back(Port("float32", PortType::FLOAT));

    // Expandable Node 4 (ResNet-style Add block)
    Node node4;
    node4.id = nextNodeId_++;
    node4.name = "Residual Add";
    node4.isExpandable = true;
    node4.parentId = -1;
    node4.inputs.push_back(Port("branch_a float32", PortType::FLOAT));
    node4.inputs.push_back(Port("branch_b float32", PortType::FLOAT));
    node4.outputs.push_back(Port("float32", PortType::FLOAT));

    // Children of Node 4 (Simple add operation internals)
    Node node4_child0;
    node4_child0.id = nextNodeId_++;
    node4_child0.name = "Add";
    node4_child0.isExpandable = false;
    node4_child0.parentId = node4.id;
    node4_child0.inputs.push_back(Port("a float32", PortType::FLOAT));
    node4_child0.inputs.push_back(Port("b float32", PortType::FLOAT));
    node4_child0.outputs.push_back(Port("float32", PortType::FLOAT));
    node4_child0.parentInputPortMap = { 0, 1 }; // Maps to both parent inputs
    node4_child0.parentOutputPortMap = { -1 }; // Internal

    Node node4_child1;
    node4_child1.id = nextNodeId_++;
    node4_child1.name = "Scale [0.5]";
    node4_child1.isExpandable = false;
    node4_child1.parentId = node4.id;
    node4_child1.inputs.push_back(Port("float32", PortType::FLOAT));
    node4_child1.outputs.push_back(Port("float32", PortType::FLOAT));
    node4_child1.parentInputPortMap = { -1 }; // Internal
    node4_child1.parentOutputPortMap = { 0 }; // Maps to parent output

    node4.childNodeIds.push_back(node4_child0.id);
    node4.childNodeIds.push_back(node4_child1.id);

    // Regular Node 5
    Node node5;
    node5.id = nextNodeId_++;
    node5.name = "Flatten";
    node5.isExpandable = false;
    node5.parentId = -1;
    node5.inputs.push_back(Port("float32", PortType::FLOAT));
    node5.outputs.push_back(Port("float32", PortType::FLOAT));

    // Regular Node 6
    Node node6;
    node6.id = nextNodeId_++;
    node6.name = "Dense [512]";
    node6.isExpandable = false;
    node6.parentId = -1;
    node6.inputs.push_back(Port("float32", PortType::FLOAT));
    node6.outputs.push_back(Port("float32", PortType::FLOAT));

    // Regular Node 7
    Node node7;
    node7.id = nextNodeId_++;
    node7.name = "Dense [256]";
    node7.isExpandable = false;
    node7.parentId = -1;
    node7.inputs.push_back(Port("float32", PortType::FLOAT));
    node7.outputs.push_back(Port("float32", PortType::FLOAT));

    // Expandable Node 8 (Attention mechanism - complex multi-input/output)
    Node node8;
    node8.id = nextNodeId_++;
    node8.name = "Multi-Head Attention";
    node8.isExpandable = true;
    node8.parentId = -1;
    node8.inputs.push_back(Port("query float32", PortType::FLOAT));
    node8.inputs.push_back(Port("key float32", PortType::FLOAT));
    node8.inputs.push_back(Port("value float32", PortType::FLOAT));
    node8.outputs.push_back(Port("float32", PortType::FLOAT));
    node8.outputs.push_back(Port("weights float64", PortType::FLOAT));

    // Children of Node 8 (Complex attention internals with many ports)
    Node node8_child0;
    node8_child0.id = nextNodeId_++;
    node8_child0.name = "Q Linear [256]";
    node8_child0.isExpandable = false;
    node8_child0.parentId = node8.id;
    node8_child0.inputs.push_back(Port("float32", PortType::FLOAT));
    node8_child0.outputs.push_back(Port("q_proj float32", PortType::FLOAT));
    node8_child0.outputs.push_back(Port("q_bias float32", PortType::FLOAT));
    node8_child0.parentInputPortMap = { 0 };  // Maps to query input
    node8_child0.parentOutputPortMap = { -1, -1 }; // Internal

    Node node8_child1;
    node8_child1.id = nextNodeId_++;
    node8_child1.name = "K Linear [256]";
    node8_child1.isExpandable = false;
    node8_child1.parentId = node8.id;
    node8_child1.inputs.push_back(Port("float32", PortType::FLOAT));
    node8_child1.outputs.push_back(Port("k_proj float32", PortType::FLOAT));
    node8_child1.outputs.push_back(Port("k_bias float32", PortType::FLOAT));
    node8_child1.parentInputPortMap = { 1 };  // Maps to key input
    node8_child1.parentOutputPortMap = { -1, -1 }; // Internal

    Node node8_child2;
    node8_child2.id = nextNodeId_++;
    node8_child2.name = "V Linear [256]";
    node8_child2.isExpandable = false;
    node8_child2.parentId = node8.id;
    node8_child2.inputs.push_back(Port("float32", PortType::FLOAT));
    node8_child2.outputs.push_back(Port("v_proj float32", PortType::FLOAT));
    node8_child2.outputs.push_back(Port("v_bias float32", PortType::FLOAT));
    node8_child2.parentInputPortMap = { 2 };  // Maps to value input
    node8_child2.parentOutputPortMap = { -1, -1 }; // Internal

    Node node8_child3;
    node8_child3.id = nextNodeId_++;
    node8_child3.name = "MatMul [Q•K^T]";
    node8_child3.isExpandable = false;
    node8_child3.parentId = node8.id;
    node8_child3.inputs.push_back(Port("q float32", PortType::FLOAT));
    node8_child3.inputs.push_back(Port("k float32", PortType::FLOAT));
    node8_child3.outputs.push_back(Port("scores float32", PortType::FLOAT));
    node8_child3.parentInputPortMap = { -1, -1 }; // Internal
    node8_child3.parentOutputPortMap = { -1 }; // Internal

    Node node8_child4;
    node8_child4.id = nextNodeId_++;
    node8_child4.name = "Softmax";
    node8_child4.isExpandable = false;
    node8_child4.parentId = node8.id;
    node8_child4.inputs.push_back(Port("float32", PortType::FLOAT));
    node8_child4.outputs.push_back(Port("attn_weights float32", PortType::FLOAT));
    node8_child4.outputs.push_back(Port("max_value float64", PortType::FLOAT));
    node8_child4.parentInputPortMap = { -1 }; // Internal
    node8_child4.parentOutputPortMap = { -1, 1 }; // Second output maps to parent weights output

    Node node8_child5;
    node8_child5.id = nextNodeId_++;
    node8_child5.name = "MatMul [Attn•V]";
    node8_child5.isExpandable = false;
    node8_child5.parentId = node8.id;
    node8_child5.inputs.push_back(Port("attn float32", PortType::FLOAT));
    node8_child5.inputs.push_back(Port("v float32", PortType::FLOAT));
    node8_child5.outputs.push_back(Port("float32", PortType::FLOAT));
    node8_child5.parentInputPortMap = { -1, -1 }; // Internal
    node8_child5.parentOutputPortMap = { 0 }; // Maps to parent output

    node8.childNodeIds.push_back(node8_child0.id);
    node8.childNodeIds.push_back(node8_child1.id);
    node8.childNodeIds.push_back(node8_child2.id);
    node8.childNodeIds.push_back(node8_child3.id);
    node8.childNodeIds.push_back(node8_child4.id);
    node8.childNodeIds.push_back(node8_child5.id);

    // Regular Node 9
    Node node9;
    node9.id = nextNodeId_++;
    node9.name = "LayerNorm";
    node9.isExpandable = false;
    node9.parentId = -1;
    node9.inputs.push_back(Port("float32", PortType::FLOAT));
    node9.outputs.push_back(Port("float32", PortType::FLOAT));

    // Regular Node 10
    Node node10;
    node10.id = nextNodeId_++;
    node10.name = "Softmax [Logits]";
    node10.isExpandable = false;
    node10.parentId = -1;
    node10.inputs.push_back(Port("float32", PortType::FLOAT));
    node10.outputs.push_back(Port("float32", PortType::FLOAT));

    // Regular Node 11
    Node node11;
    node11.id = nextNodeId_++;
    node11.name = "Output [1000 classes]";
    node11.isExpandable = false;
    node11.parentId = -1;
    node11.inputs.push_back(Port("float32", PortType::FLOAT));

    // Add all nodes to model
    nodes_.push_back(node0);
    nodes_.push_back(node1);
    nodes_.push_back(node1_child0);
    nodes_.push_back(node1_child1);
    nodes_.push_back(node1_child2);
    nodes_.push_back(node1_child3);
    nodes_.push_back(node2);
    nodes_.push_back(node3);
    nodes_.push_back(node4);
    nodes_.push_back(node4_child0);
    nodes_.push_back(node4_child1);
    nodes_.push_back(node5);
    nodes_.push_back(node6);
    nodes_.push_back(node7);
    nodes_.push_back(node8);
    nodes_.push_back(node8_child0);
    nodes_.push_back(node8_child1);
    nodes_.push_back(node8_child2);
    nodes_.push_back(node8_child3);
    nodes_.push_back(node8_child4);
    nodes_.push_back(node8_child5);
    nodes_.push_back(node9);
    nodes_.push_back(node10);
    nodes_.push_back(node11);

    // Build index
    for (size_t i = 0; i < nodes_.size(); i++) {
        nodeIdToIndex_[nodes_[i].id] = i;
    }

    // External connections (main graph flow - typical CNN architecture)
    connections_.push_back(Connection(node0.id, 0, node1.id, 0));   // Input → Conv Block
    connections_.push_back(Connection(node1.id, 0, node2.id, 0));   // Conv Block → MaxPool
    connections_.push_back(Connection(node2.id, 0, node3.id, 0));   // MaxPool → Conv2D
    connections_.push_back(Connection(node3.id, 0, node4.id, 0));   // Conv2D → Residual (branch a)
    connections_.push_back(Connection(node2.id, 0, node4.id, 1));   // MaxPool → Residual (branch b, skip connection)
    connections_.push_back(Connection(node4.id, 0, node5.id, 0));   // Residual → Flatten
    connections_.push_back(Connection(node5.id, 0, node6.id, 0));   // Flatten → Dense[512]
    connections_.push_back(Connection(node6.id, 0, node7.id, 0));   // Dense[512] → Dense[256]
    connections_.push_back(Connection(node7.id, 0, node8.id, 0));   // Dense[256] → Attention (query)
    connections_.push_back(Connection(node7.id, 0, node8.id, 1));   // Dense[256] → Attention (key)
    connections_.push_back(Connection(node7.id, 0, node8.id, 2));   // Dense[256] → Attention (value)
    connections_.push_back(Connection(node8.id, 0, node9.id, 0));   // Attention → LayerNorm
    connections_.push_back(Connection(node9.id, 0, node10.id, 0));  // LayerNorm → Softmax
    connections_.push_back(Connection(node10.id, 0, node11.id, 0)); // Softmax → Output

    // Internal connections for Node 1 (Conv2D Block)
    connections_.push_back(Connection(node1_child0.id, 0, node1_child1.id, 0));  // Conv → BatchNorm (data)
    connections_.push_back(Connection(node1_child0.id, 1, node1_child1.id, 1));  // Conv → BatchNorm (stats)
    connections_.push_back(Connection(node1_child1.id, 0, node1_child2.id, 0));  // BatchNorm → ReLU
    connections_.push_back(Connection(node1_child2.id, 0, node1_child3.id, 0));  // ReLU → Dropout

    // Internal connections for Node 4 (Residual Add)
    connections_.push_back(Connection(node4_child0.id, 0, node4_child1.id, 0));  // Add → Scale

    // Internal connections for Node 8 (Multi-Head Attention - complex graph)
    connections_.push_back(Connection(node8_child0.id, 0, node8_child3.id, 0));  // Q → MatMul[Q•K^T]
    connections_.push_back(Connection(node8_child1.id, 0, node8_child3.id, 1));  // K → MatMul[Q•K^T]
    connections_.push_back(Connection(node8_child3.id, 0, node8_child4.id, 0));  // Scores → Softmax
    connections_.push_back(Connection(node8_child4.id, 0, node8_child5.id, 0));  // Attn weights → MatMul[Attn•V]
    connections_.push_back(Connection(node8_child2.id, 0, node8_child5.id, 1));  // V → MatMul[Attn•V]

    printf("DataModel: Created AI Model Profiler Test Hierarchy\n");
    printf("  - 12 root nodes (3 expandable: Conv2D Block, Residual Add, Multi-Head Attention)\n");
    printf("  - Conv2D Block: 4 children (Conv2D, BatchNorm, ReLU, Dropout)\n");
    printf("  - Residual Add: 2 children (Add, Scale)\n");
    printf("  - Multi-Head Attention: 6 children with complex multi-input/output graph\n");
    printf("  - Total: %zu nodes, %zu connections\n", nodes_.size(), connections_.size());
}


// Find which child provides the given parent output port
const Node* DataModel::FindChildProvidingOutputPort(int parentId, int parentPortIdx) const {
    const Node* parent = GetNode(parentId);
    if (!parent) return nullptr;

    for (int childId : parent->childNodeIds) {
        const Node* child = GetNode(childId);
        if (!child) continue;

        // Check if this child's parentOutputPortMap contains the parent port index
        for (size_t i = 0; i < child->parentOutputPortMap.size(); i++) {
            if (child->parentOutputPortMap[i] == parentPortIdx) {
                return child;  // This child provides parent output[parentPortIdx]
            }
        }
    }

    return nullptr;  // No child provides this port
}

// Find which child accepts the given parent input port
const Node* DataModel::FindChildAcceptingInputPort(int parentId, int parentPortIdx) const {
    const Node* parent = GetNode(parentId);
    if (!parent) return nullptr;

    for (int childId : parent->childNodeIds) {
        const Node* child = GetNode(childId);
        if (!child) continue;

        // Check if this child's parentInputPortMap contains the parent port index
        for (size_t i = 0; i < child->parentInputPortMap.size(); i++) {
            if (child->parentInputPortMap[i] == parentPortIdx) {
                return child;  // This child accepts parent input[parentPortIdx]
            }
        }
    }

    return nullptr;  // No child accepts this port
}

// Find which child port index corresponds to the parent port index
int DataModel::FindChildPortIndex(const std::vector<int>& portMap, int parentPortIdx) const {
    for (size_t i = 0; i < portMap.size(); i++) {
        if (portMap[i] == parentPortIdx) {
            return i;  // Child port[i] represents parent port[parentPortIdx]
        }
    }
    return -1;  // Not found
}

void DataModel::Clear() {
    nodes_.clear();
    connections_.clear();
    nodeIdToIndex_.clear();
    nextNodeId_ = 0;
}


void DataModel::CreateLargeTestHierarchy() {
    Clear();

    // Helper lambda: Create a simple node and add to graph
    auto addNode = [&](const std::string& name, bool expandable, int parentId,
                       const std::vector<std::string>& inputNames,
                       const std::vector<std::string>& outputNames,
                       const std::vector<int>& inputPortMap = {},
                       const std::vector<int>& outputPortMap = {}) -> int {
        Node node;
        node.id = nextNodeId_++;
        node.name = name;
        node.isExpandable = expandable;
        node.parentId = parentId;
        for (const auto& inName : inputNames) {
            node.inputs.push_back(Port(inName, PortType::FLOAT));
        }
        for (const auto& outName : outputNames) {
            node.outputs.push_back(Port(outName, PortType::FLOAT));
        }
        node.parentInputPortMap = inputPortMap;
        node.parentOutputPortMap = outputPortMap;
        nodes_.push_back(node);
        nodeIdToIndex_[node.id] = nodes_.size() - 1;
        return node.id;
    };

    // Helper: Connect node to previous
    auto connect = [&](int srcId, int srcPort, int dstId, int dstPort) {
        connections_.push_back(Connection(srcId, srcPort, dstId, dstPort));
    };

    // Helper: Create CNN Bottleneck Block (expandable, 6 children)
    auto createBottleneckBlock = [&](const std::string& stageName, int blockIdx, int lastNodeId) -> std::pair<int, int> {
        int blockId = addNode(stageName + "_Block" + std::to_string(blockIdx), true, -1,
                              {"x"}, {"x", "features"});

        int conv1 = addNode("Conv1x1 Reduce", false, blockId, {"x"}, {"reduced"}, {0}, {-1});
        int bn1 = addNode("BatchNorm", false, blockId, {"reduced"}, {"norm"}, {-1}, {-1});
        int relu1 = addNode("ReLU", false, blockId, {"norm"}, {"act"}, {-1}, {-1});
        int conv3 = addNode("Conv3x3", false, blockId, {"act"}, {"features"}, {-1}, {1});
        int bn2 = addNode("BatchNorm", false, blockId, {"features"}, {"norm"}, {-1}, {-1});
        int relu2 = addNode("ReLU", false, blockId, {"norm"}, {"act"}, {-1}, {-1});
        int conv1_exp = addNode("Conv1x1 Expand", false, blockId, {"act"}, {"expanded"}, {-1}, {-1});
        int bn3 = addNode("BatchNorm", false, blockId, {"expanded"}, {"norm"}, {-1}, {-1});
        int add = addNode("Residual Add", false, blockId, {"identity", "transformed"}, {"sum"}, {0, -1}, {-1});
        int relu_out = addNode("ReLU", false, blockId, {"sum"}, {"out"}, {-1}, {0});

        Node* block = GetNodeMutable(blockId);
        block->childNodeIds = {conv1, bn1, relu1, conv3, bn2, relu2, conv1_exp, bn3, add, relu_out};

        connect(conv1, 0, bn1, 0);
        connect(bn1, 0, relu1, 0);
        connect(relu1, 0, conv3, 0);
        connect(conv3, 0, bn2, 0);
        connect(bn2, 0, relu2, 0);
        connect(relu2, 0, conv1_exp, 0);
        connect(conv1_exp, 0, bn3, 0);
        connect(bn3, 0, add, 1);
        connect(add, 0, relu_out, 0);

        if (lastNodeId != -1) {
            connect(lastNodeId, 0, blockId, 0);
        }

        return {blockId, blockId}; // {main_out, features_out}
    };

    // Helper: Create Transformer Block (expandable, ~18 children with 2-level hierarchy)
    auto createTransformerBlock = [&](const std::string& stageName, int blockIdx, int lastNodeId) -> int {
        int blockId = addNode(stageName + "_Block" + std::to_string(blockIdx), true, -1,
                              {"x"}, {"x"});

        // Attention path
        int norm1 = addNode("LayerNorm", false, blockId, {"x"}, {"norm"}, {0}, {-1});

        // Multi-Head Attention (expandable sub-block)
        int attnId = addNode("Multi-Head Attention", true, blockId, {"x"}, {"attn_out"}, {-1}, {-1});
        int q_proj = addNode("Q Projection", false, attnId, {"x"}, {"Q"}, {0}, {-1});
        int k_proj = addNode("K Projection", false, attnId, {"x"}, {"K"}, {0}, {-1});
        int v_proj = addNode("V Projection", false, attnId, {"x"}, {"V"}, {0}, {-1});
        int attn_scores = addNode("Attention Scores", false, attnId, {"Q", "K"}, {"scores"}, {-1, -1}, {-1});
        int softmax = addNode("Softmax", false, attnId, {"scores"}, {"weights"}, {-1}, {-1});
        int attn_drop = addNode("Dropout", false, attnId, {"weights"}, {"weights"}, {-1}, {-1});
        int attn_mm = addNode("MatMul [W•V]", false, attnId, {"weights", "V"}, {"context"}, {-1, -1}, {-1});
        int out_proj = addNode("Output Projection", false, attnId, {"context"}, {"out"}, {-1}, {0});

        Node* attn = GetNodeMutable(attnId);
        attn->childNodeIds = {q_proj, k_proj, v_proj, attn_scores, softmax, attn_drop, attn_mm, out_proj};

        connect(q_proj, 0, attn_scores, 0);
        connect(k_proj, 0, attn_scores, 1);
        connect(attn_scores, 0, softmax, 0);
        connect(softmax, 0, attn_drop, 0);
        connect(attn_drop, 0, attn_mm, 0);
        connect(v_proj, 0, attn_mm, 1);
        connect(attn_mm, 0, out_proj, 0);

        int attn_drop_out = addNode("Dropout", false, blockId, {"attn"}, {"attn"}, {-1}, {-1});
        int res1 = addNode("Residual Add", false, blockId, {"x", "attn"}, {"x"}, {0, -1}, {-1});

        // FFN path
        int norm2 = addNode("LayerNorm", false, blockId, {"x"}, {"norm"}, {-1}, {-1});

        // Feed-Forward Network (expandable sub-block)
        int ffnId = addNode("Feed-Forward", true, blockId, {"x"}, {"out"}, {-1}, {-1});
        int fc1 = addNode("Linear [expand 4x]", false, ffnId, {"x"}, {"hidden"}, {0}, {-1});
        int gelu = addNode("GELU", false, ffnId, {"hidden"}, {"act"}, {-1}, {-1});
        int drop1 = addNode("Dropout", false, ffnId, {"act"}, {"hidden"}, {-1}, {-1});
        int fc2 = addNode("Linear [project]", false, ffnId, {"hidden"}, {"out"}, {-1}, {0});

        Node* ffn = GetNodeMutable(ffnId);
        ffn->childNodeIds = {fc1, gelu, drop1, fc2};

        connect(fc1, 0, gelu, 0);
        connect(gelu, 0, drop1, 0);
        connect(drop1, 0, fc2, 0);

        int ffn_drop_out = addNode("Dropout", false, blockId, {"ffn"}, {"ffn"}, {-1}, {-1});
        int res2 = addNode("Residual Add", false, blockId, {"x", "ffn"}, {"x"}, {-1, -1}, {0});

        Node* block = GetNodeMutable(blockId);
        block->childNodeIds = {norm1, attnId, attn_drop_out, res1, norm2, ffnId, ffn_drop_out, res2};

        connect(norm1, 0, attnId, 0);
        connect(attnId, 0, attn_drop_out, 0);
        connect(attn_drop_out, 0, res1, 1);
        connect(res1, 0, norm2, 0);
        connect(norm2, 0, ffnId, 0);
        connect(ffnId, 0, ffn_drop_out, 0);
        connect(ffn_drop_out, 0, res2, 1);

        if (lastNodeId != -1) {
            connect(lastNodeId, 0, blockId, 0);
        }

        return blockId;
    };

    // ============================================================================
    // HYBRID CNN-TRANSFORMER ARCHITECTURE - ~1000 nodes
    // ============================================================================
    // Architecture inspired by CoAtNet / Swin Transformer
    // Structure:
    //   Stage 1 (CNN): Input → Stem → 3 CNN Bottleneck Blocks
    //   Stage 2 (CNN): 4 CNN Bottleneck Blocks → Aux Classifier 1
    //   Stage 3 (Hybrid): 6 CNN Blocks → Transition → 4 Transformer Blocks
    //   Stage 4 (Transformer): 12 Transformer Blocks → Aux Classifier 2
    //   Stage 5 (Transformer): 18 Transformer Blocks
    //   FPN: Multi-scale feature fusion (P2/P3/P4/P5)
    //   Heads: Classification, Segmentation, Detection (parallel outputs)
    // Features:
    //   ✓ Parallel branches (FPN + multiple task heads)
    //   ✓ Long-range skip connections (ResNet residuals + FPN laterals)
    //   ✓ Multiple outputs (3 task heads + 2 auxiliary classifiers)
    //   ✓ 2-level hierarchy (blocks contain sub-blocks like Attention/FFN)
    // ============================================================================

    printf("\n=== Creating Hybrid CNN-Transformer Model (~1000 nodes) ===\n");

    // ----------------------------------------------------------------------------
    // INPUT STEM (3 nodes)
    // ----------------------------------------------------------------------------
    int input = addNode("Input [224x224x3]", false, -1, {}, {"image"});
    int stem_conv = addNode("Stem Conv 7x7/2", false, -1, {"image"}, {"features"});
    int stem_pool = addNode("MaxPool 3x3/2", false, -1, {"features"}, {"pooled"});

    connect(input, 0, stem_conv, 0);
    connect(stem_conv, 0, stem_pool, 0);

    int lastNode = stem_pool;
    std::vector<int> stage1_features, stage2_features, stage3_features, stage4_features, stage5_features;

    printf("  Stem: 3 nodes\n");

    // ----------------------------------------------------------------------------
    // STAGE 1: CNN Backbone (3 bottleneck blocks × 10 nodes each = 30 nodes)
    // ----------------------------------------------------------------------------
    for (int i = 0; i < 3; i++) {
        auto [blockId, featId] = createBottleneckBlock("Stage1", i, lastNode);
        lastNode = blockId;
        stage1_features.push_back(featId);
    }
    printf("  Stage 1 (CNN): 3 blocks × 11 nodes = 33 nodes\n");

    // ----------------------------------------------------------------------------
    // STAGE 2: CNN Backbone (4 bottleneck blocks = 40 nodes) + Auxiliary Classifier
    // ----------------------------------------------------------------------------
    for (int i = 0; i < 4; i++) {
        auto [blockId, featId] = createBottleneckBlock("Stage2", i, lastNode);
        lastNode = blockId;
        stage2_features.push_back(featId);
    }

    // Auxiliary Classifier 1 (early exit)
    int aux1 = addNode("Aux Classifier 1", true, -1, {"features"}, {"logits [1000]"});
    int aux1_pool = addNode("AdaptiveAvgPool", false, aux1, {"features"}, {"pooled"}, {0}, {-1});
    int aux1_fc = addNode("FC [512→1000]", false, aux1, {"pooled"}, {"logits"}, {-1}, {0});
    GetNodeMutable(aux1)->childNodeIds = {aux1_pool, aux1_fc};
    connect(aux1_pool, 0, aux1_fc, 0);
    connect(lastNode, 1, aux1, 0); // Connect to features output

    printf("  Stage 2 (CNN): 4 blocks × 11 nodes + Aux1 (3 nodes) = 47 nodes\n");

    // ----------------------------------------------------------------------------
    // STAGE 3: Hybrid (6 CNN blocks + 4 Transformer blocks = 178 nodes)
    // ----------------------------------------------------------------------------
    for (int i = 0; i < 6; i++) {
        auto [blockId, featId] = createBottleneckBlock("Stage3_CNN", i, lastNode);
        lastNode = blockId;
        stage3_features.push_back(featId);
    }

    // Transition: Reshape for transformers
    int transition = addNode("CNN→Transformer Transition", false, -1, {"cnn_features"}, {"tokens"});
    connect(lastNode, 0, transition, 0);
    lastNode = transition;

    for (int i = 0; i < 4; i++) {
        lastNode = createTransformerBlock("Stage3_Transformer", i, lastNode);
        stage3_features.push_back(lastNode);
    }

    printf("  Stage 3 (Hybrid): 6 CNN blocks (66 nodes) + transition (1 node) + 4 Transformer blocks (76 nodes) = 143 nodes\n");

    // ----------------------------------------------------------------------------
    // STAGE 4: Transformer (12 blocks = 228 nodes) + Auxiliary Classifier 2
    // ----------------------------------------------------------------------------
    for (int i = 0; i < 12; i++) {
        lastNode = createTransformerBlock("Stage4", i, lastNode);
        stage4_features.push_back(lastNode);
    }

    // Auxiliary Classifier 2
    int aux2 = addNode("Aux Classifier 2", true, -1, {"features"}, {"logits [1000]"});
    int aux2_pool = addNode("AdaptiveAvgPool", false, aux2, {"features"}, {"pooled"}, {0}, {-1});
    int aux2_fc = addNode("FC [1024→1000]", false, aux2, {"pooled"}, {"logits"}, {-1}, {0});
    GetNodeMutable(aux2)->childNodeIds = {aux2_pool, aux2_fc};
    connect(aux2_pool, 0, aux2_fc, 0);
    connect(lastNode, 0, aux2, 0);

    printf("  Stage 4 (Transformer): 12 blocks × 19 nodes + Aux2 (3 nodes) = 231 nodes\n");

    // ----------------------------------------------------------------------------
    // STAGE 5: Deep Transformer (18 blocks = 342 nodes)
    // ----------------------------------------------------------------------------
    for (int i = 0; i < 18; i++) {
        lastNode = createTransformerBlock("Stage5", i, lastNode);
        stage5_features.push_back(lastNode);
    }

    printf("  Stage 5 (Transformer): 18 blocks × 19 nodes = 342 nodes\n");

    // ----------------------------------------------------------------------------
    // FEATURE PYRAMID NETWORK (FPN) - Multi-scale feature fusion (20 nodes)
    // ----------------------------------------------------------------------------
    // Use features from different stages for multi-scale processing
    int fpn_p5 = addNode("FPN_P5 [7x7]", false, -1, {"C5"}, {"P5"});
    connect(stage5_features.back(), 0, fpn_p5, 0);

    int fpn_up_p5 = addNode("Upsample 2x", false, -1, {"P5"}, {"P5_up"});
    int fpn_lat_c4 = addNode("Lateral Conv1x1", false, -1, {"C4"}, {"C4_lateral"});
    int fpn_p4 = addNode("FPN_P4 Merge [14x14]", false, -1, {"up", "lateral"}, {"P4"});
    connect(fpn_p5, 0, fpn_up_p5, 0);
    connect(stage4_features.back(), 0, fpn_lat_c4, 0);
    connect(fpn_up_p5, 0, fpn_p4, 0);
    connect(fpn_lat_c4, 0, fpn_p4, 1);

    int fpn_up_p4 = addNode("Upsample 2x", false, -1, {"P4"}, {"P4_up"});
    int fpn_lat_c3 = addNode("Lateral Conv1x1", false, -1, {"C3"}, {"C3_lateral"});
    int fpn_p3 = addNode("FPN_P3 Merge [28x28]", false, -1, {"up", "lateral"}, {"P3"});
    connect(fpn_p4, 0, fpn_up_p4, 0);
    connect(stage3_features.back(), 0, fpn_lat_c3, 0);
    connect(fpn_up_p4, 0, fpn_p3, 0);
    connect(fpn_lat_c3, 0, fpn_p3, 1);

    int fpn_up_p3 = addNode("Upsample 2x", false, -1, {"P3"}, {"P3_up"});
    int fpn_lat_c2 = addNode("Lateral Conv1x1", false, -1, {"C2"}, {"C2_lateral"});
    int fpn_p2 = addNode("FPN_P2 Merge [56x56]", false, -1, {"up", "lateral"}, {"P2"});
    connect(fpn_p3, 0, fpn_up_p3, 0);
    connect(stage2_features.back(), 0, fpn_lat_c2, 0);
    connect(fpn_up_p3, 0, fpn_p2, 0);
    connect(fpn_lat_c2, 0, fpn_p2, 1);

    printf("  FPN: 13 nodes (4-scale pyramid with lateral connections)\n");

    // ----------------------------------------------------------------------------
    // TASK HEADS - Parallel Multi-Task Outputs (30 nodes total)
    // ----------------------------------------------------------------------------

    // === Classification Head (5 nodes) ===
    int cls_head = addNode("Classification Head", true, -1, {"features"}, {"class_logits [1000]"});
    int cls_pool = addNode("Global AvgPool", false, cls_head, {"features"}, {"pooled"}, {0}, {-1});
    int cls_drop = addNode("Dropout [0.5]", false, cls_head, {"pooled"}, {"dropped"}, {-1}, {-1});
    int cls_fc = addNode("FC [2048→1000]", false, cls_head, {"dropped"}, {"logits"}, {-1}, {0});
    int cls_softmax = addNode("Softmax", false, -1, {"logits"}, {"probabilities"});

    GetNodeMutable(cls_head)->childNodeIds = {cls_pool, cls_drop, cls_fc};
    connect(cls_pool, 0, cls_drop, 0);
    connect(cls_drop, 0, cls_fc, 0);
    connect(fpn_p5, 0, cls_head, 0);
    connect(cls_head, 0, cls_softmax, 0);

    // === Segmentation Head (9 nodes) ===
    int seg_head = addNode("Segmentation Head", true, -1, {"P2", "P3", "P4"}, {"seg_map [224x224x21]"});
    int seg_up_p4 = addNode("Upsample P4→56x56", false, seg_head, {"P4"}, {"P4_up"}, {2}, {-1});
    int seg_up_p3 = addNode("Upsample P3→56x56", false, seg_head, {"P3"}, {"P3_up"}, {1}, {-1});
    int seg_concat = addNode("Concat Multi-Scale", false, seg_head, {"P2", "P3_up", "P4_up"}, {"concat"}, {0, -1, -1}, {-1});
    int seg_conv = addNode("Conv 3x3 [C→256]", false, seg_head, {"concat"}, {"features"}, {-1}, {-1});
    int seg_up_final = addNode("Upsample→224x224", false, seg_head, {"features"}, {"upsampled"}, {-1}, {-1});
    int seg_classifier = addNode("Conv 1x1 [256→21]", false, seg_head, {"upsampled"}, {"seg_map"}, {-1}, {0});

    GetNodeMutable(seg_head)->childNodeIds = {seg_up_p4, seg_up_p3, seg_concat, seg_conv, seg_up_final, seg_classifier};
    connect(seg_up_p4, 0, seg_concat, 2);
    connect(seg_up_p3, 0, seg_concat, 1);
    connect(seg_concat, 0, seg_conv, 0);
    connect(seg_conv, 0, seg_up_final, 0);
    connect(seg_up_final, 0, seg_classifier, 0);
    connect(fpn_p2, 0, seg_head, 0);
    connect(fpn_p3, 0, seg_head, 1);
    connect(fpn_p4, 0, seg_head, 2);

    // === Detection Head (9 nodes) ===
    int det_head = addNode("Detection Head", true, -1, {"P3", "P4", "P5"}, {"bbox", "class"});
    int det_rpn_p3 = addNode("RPN P3", false, det_head, {"P3"}, {"proposals_p3"}, {0}, {-1});
    int det_rpn_p4 = addNode("RPN P4", false, det_head, {"P4"}, {"proposals_p4"}, {1}, {-1});
    int det_rpn_p5 = addNode("RPN P5", false, det_head, {"P5"}, {"proposals_p5"}, {2}, {-1});
    int det_roi = addNode("ROI Align", false, det_head, {"proposals"}, {"roi_features"}, {-1}, {-1});
    int det_bbox = addNode("BBox Regressor", false, det_head, {"roi_features"}, {"bbox_deltas"}, {-1}, {0});
    int det_class = addNode("Class Predictor [80]", false, det_head, {"roi_features"}, {"class_scores"}, {-1}, {1});

    GetNodeMutable(det_head)->childNodeIds = {det_rpn_p3, det_rpn_p4, det_rpn_p5, det_roi, det_bbox, det_class};
    connect(det_rpn_p3, 0, det_roi, 0);
    connect(det_roi, 0, det_bbox, 0);
    connect(det_roi, 0, det_class, 0);
    connect(fpn_p3, 0, det_head, 0);
    connect(fpn_p4, 0, det_head, 1);
    connect(fpn_p5, 0, det_head, 2);

    // ----------------------------------------------------------------------------
    // FINAL OUTPUT NODES - All task outputs feed into loss computations
    // ----------------------------------------------------------------------------
    int cls_loss = addNode("Classification Loss", false, -1, {"predictions", "labels"}, {"cls_loss"});
    int seg_loss = addNode("Segmentation Loss", false, -1, {"seg_map", "targets"}, {"seg_loss"});
    int det_loss = addNode("Detection Loss", false, -1, {"bbox", "class", "gt_boxes"}, {"det_loss"});
    int aux1_loss = addNode("Aux1 Loss", false, -1, {"logits", "labels"}, {"aux1_loss"});
    int aux2_loss = addNode("Aux2 Loss", false, -1, {"logits", "labels"}, {"aux2_loss"});

    // Final loss aggregation
    int total_loss = addNode("Total Loss", false, -1, {"cls", "seg", "det", "aux1", "aux2"}, {"loss"});

    connect(cls_softmax, 0, cls_loss, 0);
    connect(seg_head, 0, seg_loss, 0);
    connect(det_head, 0, det_loss, 0);
    connect(det_head, 1, det_loss, 1);
    connect(aux1, 0, aux1_loss, 0);
    connect(aux2, 0, aux2_loss, 0);

    connect(cls_loss, 0, total_loss, 0);
    connect(seg_loss, 0, total_loss, 1);
    connect(det_loss, 0, total_loss, 2);
    connect(aux1_loss, 0, total_loss, 3);
    connect(aux2_loss, 0, total_loss, 4);

    printf("  Task Heads: Classification (5 nodes) + Segmentation (7 nodes) + Detection (7 nodes) = 19 nodes\n");
    printf("  Loss Nodes: 6 nodes (5 task losses + 1 total loss)\n");

    // ----------------------------------------------------------------------------
    // SUMMARY
    // ----------------------------------------------------------------------------
    printf("\n=== Hybrid CNN-Transformer Model Created ===\n");
    printf("  Total Nodes: %zu\n", nodes_.size());
    printf("  Total Connections: %zu\n", connections_.size());
    printf("  \n");
    printf("  Architecture Summary:\n");
    printf("    - Stage 1: 3 CNN blocks (33 nodes)\n");
    printf("    - Stage 2: 4 CNN blocks + Aux1 (47 nodes)\n");
    printf("    - Stage 3: 6 CNN + 4 Transformer (143 nodes)\n");
    printf("    - Stage 4: 12 Transformer + Aux2 (231 nodes)\n");
    printf("    - Stage 5: 18 Transformer (342 nodes)\n");
    printf("    - FPN: Multi-scale fusion (13 nodes)\n");
    printf("    - Task Heads: 3 parallel outputs (19 nodes)\n");
    printf("  \n");
    printf("  Features:\n");
    printf("    ✓ Parallel branches (FPN 4-scale + 3 task heads)\n");
    printf("    ✓ Long-range skip connections (CNN residuals + FPN laterals)\n");
    printf("    ✓ Multiple outputs (Classification + Segmentation + Detection + 2 Aux)\n");
    printf("    ✓ 2-level hierarchy (Transformer blocks contain Attention + FFN subgraphs)\n");
    printf("========================================\n\n");
}

void DataModel::CreateLongEdgeTest() {
    Clear();

    // Minimal test for edge-polyline subsampling (MAX_EDGE_POINTS cap in the GPU edge path).
    // A CHAIN of N nodes forces N layers (each node one layer deeper than the last). A single
    // SKIP edge from the first node to the last then spans all N layers, so its route picks up
    // ~N dummy waypoints. With N > MAX_EDGE_POINTS (32), the route exceeds the cap and must be
    // subsampled while still reaching both true endpoints. Watch the skip edge in edge render
    // modes 1/2/3: it must connect Start -> End across the whole graph, not stop early.
    const int CHAIN_LEN = 40;  // > 32 so subsampling is exercised

    std::vector<int> chain;
    chain.reserve(CHAIN_LEN);
    for (int i = 0; i < CHAIN_LEN; i++) {
        std::string name = (i == 0) ? "Start"
                         : (i == CHAIN_LEN - 1) ? "End"
                         : ("Layer" + std::to_string(i));
        Node node;
        node.id = nextNodeId_++;
        node.name = name;
        node.isExpandable = false;
        node.parentId = -1;
        node.inputs.push_back(Port("in", PortType::FLOAT));
        node.outputs.push_back(Port("out", PortType::FLOAT));
        nodes_.push_back(node);
        nodeIdToIndex_[node.id] = nodes_.size() - 1;
        chain.push_back(node.id);
    }

    // Chain edges: establish the layers (short adjacent-layer edges).
    for (int i = 0; i + 1 < CHAIN_LEN; i++) {
        connections_.push_back(Connection(chain[i], 0, chain[i + 1], 0));
    }

    // The long skip edge under test: Start -> End, spanning every layer.
    // End's input port 0 is used by the chain; add a second input port for the skip.
    GetNodeMutable(chain[CHAIN_LEN - 1])->inputs.push_back(Port("skip", PortType::FLOAT));
    connections_.push_back(Connection(chain[0], 0, chain[CHAIN_LEN - 1], 1));

    printf("\n=== CreateLongEdgeTest ===\n");
    printf("  %d-node chain + 1 skip edge (Start -> End) spanning all layers\n", CHAIN_LEN);
    printf("  Skip route will exceed MAX_EDGE_POINTS (32) -> exercises subsampling\n");
    printf("==========================\n\n");
}

void DataModel::CreateLargeWideTestHierarchy() {
    Clear();

    printf("\n=== Creating Wide Test Graph (~1000 nodes) ===\n");

    // Helper lambda: Create a simple node
    auto addNode = [&](const std::string& name, bool expandable, int parentId,
                       const std::vector<std::string>& inputNames,
                       const std::vector<std::string>& outputNames,
                       const std::vector<int>& inputPortMap = {},
                       const std::vector<int>& outputPortMap = {}) -> int {
        Node node;
        node.id = nextNodeId_++;
        node.name = name;
        node.isExpandable = expandable;
        node.parentId = parentId;
        for (const auto& inName : inputNames) {
            node.inputs.push_back(Port(inName, PortType::FLOAT));
        }
        for (const auto& outName : outputNames) {
            node.outputs.push_back(Port(outName, PortType::FLOAT));
        }
        node.parentInputPortMap = inputPortMap;
        node.parentOutputPortMap = outputPortMap;
        nodes_.push_back(node);
        nodeIdToIndex_[node.id] = nodes_.size() - 1;
        return node.id;
    };

    auto connect = [&](int srcId, int srcPort, int dstId, int dstPort) {
        connections_.push_back(Connection(srcId, srcPort, dstId, dstPort));
    };

    // Strategy: Create a wide, shallow network with many parallel branches
    // This creates many nodes per layer for better SPH testing

    // Layer 0: Input
    int input = addNode("Input [224x224x3]", false, -1, {}, {"image"});

    // Layer 1: Initial processing (10 parallel stems)
    std::vector<int> stems;
    for (int i = 0; i < 10; i++) {
        int stem = addNode("Stem_" + std::to_string(i), false, -1, {"x"}, {"features"});
        connect(input, 0, stem, 0);
        stems.push_back(stem);
    }

    // Create 8 processing stages, each with 10-15 parallel paths
    std::vector<int> currentLayer = stems;
    std::vector<std::vector<int>> allStages;  // Store all stage outputs for skip connections
    allStages.push_back(stems);
    int totalNodes = 1 + stems.size();  // Input + stems

    for (int stage = 0; stage < 8; stage++) {
        std::vector<int> nextLayer;
        int branchesThisStage = 10 + (stage % 6);  // 10-15 branches

        for (int branch = 0; branch < branchesThisStage; branch++) {
            // Create an expandable processing block with 8 children
            std::string blockName = "Stage" + std::to_string(stage) + "_Branch" + std::to_string(branch);
            int blockId = addNode(blockName, true, -1, {"x"}, {"out"});

            // Create 8 internal nodes (simple pipeline, no sub-subgraphs)
            int conv1 = addNode("Conv1", false, blockId, {"x"}, {"f1"}, {0}, {-1});
            int bn1 = addNode("BN1", false, blockId, {"f1"}, {"f2"}, {-1}, {-1});
            int relu1 = addNode("ReLU1", false, blockId, {"f2"}, {"f3"}, {-1}, {-1});
            int conv2 = addNode("Conv2", false, blockId, {"f3"}, {"f4"}, {-1}, {-1});
            int bn2 = addNode("BN2", false, blockId, {"f4"}, {"f5"}, {-1}, {-1});
            int relu2 = addNode("ReLU2", false, blockId, {"f5"}, {"f6"}, {-1}, {-1});
            int conv3 = addNode("Conv3", false, blockId, {"f6"}, {"f7"}, {-1}, {-1});
            int bn3 = addNode("BN3", false, blockId, {"f7"}, {"out"}, {-1}, {0});

            Node* block = GetNodeMutable(blockId);
            block->childNodeIds = {conv1, bn1, relu1, conv2, bn2, relu2, conv3, bn3};

            // Connect internal pipeline
            connect(conv1, 0, bn1, 0);
            connect(bn1, 0, relu1, 0);
            connect(relu1, 0, conv2, 0);
            connect(conv2, 0, bn2, 0);
            connect(bn2, 0, relu2, 0);
            connect(relu2, 0, conv3, 0);
            connect(conv3, 0, bn3, 0);

            // Connect to previous layer (fan-in from 1-2 nodes)
            int numConnections = 1 + (branch % 2);  // 1 or 2 connections
            for (int c = 0; c < numConnections && !currentLayer.empty(); c++) {
                int srcIdx = (branch + c) % currentLayer.size();
                connect(currentLayer[srcIdx], 0, blockId, 0);
            }

            // Add skip connections (DenseNet-style) - connect to earlier stages
            if (stage >= 2 && branch % 3 == 0) {
                // Every 3rd branch gets a long-range skip connection from 2-4 stages back
                int skipDistance = 2 + (branch % 3);  // Skip 2-4 stages
                if (stage >= skipDistance) {
                    int skipStage = stage - skipDistance;
                    int skipIdx = branch % allStages[skipStage].size();
                    connect(allStages[skipStage][skipIdx], 0, blockId, 0);
                }
            }

            nextLayer.push_back(blockId);
            totalNodes += 9;  // block + 8 children
        }

        allStages.push_back(nextLayer);
        currentLayer = nextLayer;
    }

    // Final layer: Merge to output heads (5 parallel heads)
    std::vector<int> heads;
    for (int i = 0; i < 5; i++) {
        int head = addNode("Head_" + std::to_string(i), false, -1, {"features"}, {"output"});

        // Each head connects to multiple nodes from last layer
        for (int j = 0; j < 3 && j < (int)currentLayer.size(); j++) {
            int srcIdx = (i * 3 + j) % currentLayer.size();
            connect(currentLayer[srcIdx], 0, head, 0);
        }

        heads.push_back(head);
        totalNodes++;
    }

    printf("  Total Nodes: %d\n", totalNodes);
    printf("  Total Connections: %zu\n", connections_.size());
    printf("  Structure:\n");
    printf("    - 10 parallel stems\n");
    printf("    - 8 stages with 10-15 parallel branches each\n");
    printf("    - Each branch: expandable block with 8 children\n");
    printf("    - 5 output heads\n");
    printf("  \n");
    printf("  Features:\n");
    printf("    ✓ Wide layout: 10-15 nodes per layer\n");
    printf("    ✓ No nested subgraphs (only 1-level hierarchy)\n");
    printf("    ✓ Skip connections spanning 2-4 stages (tests long edges)\n");
    printf("    ✓ All nodes have outputs (realistic AI model structure)\n");
    printf("    ✓ Good for SPH testing (nodes surrounded above/below)\n");
    printf("========================================\n\n");
}


} // namespace GraphData