<div id="graph-container_${unique_id}">
    <div id="graph_${unique_id}"></div>
    <div id="popup_${unique_id}" class="popup_${unique_id}" style="display: none;">
        <div class="popup-header_${unique_id}">
            <h3 id="popup-title_${unique_id}"></h3>
            <span class="popup-close_${unique_id}">✕</span>
        </div>
        <div id="popup-content_${unique_id}" class="popup-content_${unique_id}">
            <div id="common-info_${unique_id}">
            </div>
            <div id="module-template_${unique_id}" style="display:none;">
                <details open id="module-parameters-details_${unique_id}">
                    <summary style="color: black"><strong>forward() Parameters</strong></summary>
                    <div id="module-parameters_${unique_id}" class="json-viewer_${unique_id}"></div>
                </details>
                <details open id="module-attributes-details_${unique_id}">
                    <summary style="color: black"><strong>Attributes</strong></summary>
                    <div id="module-attributes_${unique_id}" class="json-viewer_${unique_id}"></div>
                </details>
            </div>
            <div id="tensor-op-template_${unique_id}" style="display:none;">
                <summary><strong>Parameters</strong></summary>
                <div id="tensor-op-parameters_${unique_id}" class="json-viewer_${unique_id}"></div>
            </div>
            <p id="no-info_${unique_id}" style="display:none;">No further information available for this node.</p>
        </div>
    </div>

    <style>
        .node_${unique_id} rect {
            stroke: #333;
            fill: #fff;
        }
        text.node-label_${unique_id} {
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
            cursor: default;
            user-select: none;
        }
        text.type-indicator_${unique_id} {
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
            cursor: default;
            user-select: none;
        }
        text.collapse-icon_${unique_id} {
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
            font-size: 14px
            cursor: default;
            user-select: none;
        }
        .edge_${unique_id} {
            stroke: #797979;
            fill: none;
        }
        .node_${unique_id} rect.input_${unique_id} { fill: #FF9AFF; }
        .node_${unique_id} rect.constant_${unique_id} { fill: rgb(235 235 235); }

        .node_${unique_id} rect.output_${unique_id} { fill: #74e071; }
        .node_${unique_id} rect.failed_${unique_id} { fill: rgb(224 15 15); }
        .node_${unique_id} rect.module_${unique_id} { fill: #8BDEDB; }
        .node_${unique_id} rect.operation_${unique_id} { fill: #FEBC5A; }
        .edgeLabelGroup_${unique_id} {
            pointer-events: none;
        }
        #arrowhead_${unique_id} path {
            fill: #333333;;
        }
        .node_${unique_id} rect,
        .node_${unique_id} ellipse {
            transition: fill 0.3s ease;
        }

        .node_${unique_id} .type-indicator_${unique_id} {
            font-size: 5px;
            fill: rgba(0, 0, 0, 1.0);  /* Dark grey, semi-transparent */
            pointer-events: none;
        }

        .node_${unique_id}.error_${unique_id} rect {
            stroke: black;
            stroke-width: 1px;
        }

        .boundary-box_${unique_id}.error_${unique_id} > rect:first-of-type {
            fill: red;
            fill-opacity: 0.2;
        }

        #graph-container_${unique_id} {
            position: relative;
            height: ${height};
            overflow: hidden;
        }

        #graph_${unique_id} {
            height: 100%;
        }
        
        .popup_${unique_id} {
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
            position: absolute;
            background: white;
            border: 1px solid #ccc;
            width: 500px;
            resize: horizontal;
            overflow: auto;
            max-height: 70%;
            box-shadow: 0 2px 10px rgba(0,0,0,0.2);
            display: none;
            display: flex;
            flex-direction: column;
        }

        .popup-header_${unique_id} {
            position: sticky;
            top: 0;
            background: white;
            padding: 5px 10px;
            border-bottom: 1px solid #ccc;
            display: flex;
            justify-content: space-between;
            align-items: center;
            z-index: 1;
        }

        .popup_${unique_id} h3 {
            margin: 0;
            font-size: 16px;
        }

        .popup-content_${unique_id} {
            padding: 10px;
            padding-bottom: 10px;
            overflow-y: auto;
            overflow-x: auto;
            flex-grow: 1;
        }

        .popup_${unique_id} pre {
            white-space: pre;
            overflow-x: auto;
        }

        .popup-close_${unique_id} {
            cursor: pointer;
            font-size: 24px;
            font-weight: bold;
            line-height: 1;
        }

        #popup-title_${unique_id} {
            font-size: 16px;
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
            color: black;
        }

        #graph-container_${unique_id} {
            position: relative;
        }

        #graph-container_${unique_id} pre {
            margin-bottom: 0;
        }

        .boundary-box_${unique_id} rect {
            pointer-events: none;
        }
        .boundary-box_${unique_id} text {
            pointer-events: none;
            user-select: none;
        }

        .boundary-box_${unique_id} rect {
            pointer-events: none;
        }
        .boundary-box_${unique_id} text {
            pointer-events: none;
            user-select: none;
            font-weight: bold;
        }
        .edge-label_${unique_id} text {
            pointer-events: none;
            user-select: none;
            font-size: 4px;
            font-family: "'Roboto', sans-serif";
        }
                    
        .collapse-button_${unique_id}, .info-button_${unique_id} {
            fill: #fff;
            stroke: #333;
            stroke-width: 1px;
            cursor: pointer;
        }

        .collapse-button_${unique_id}:hover, .info-button_${unique_id}:hover {
            fill: #eee;
            cursor: pointer;
        }

        .collapse-icon_${unique_id}, .info-icon_${unique_id} {
            pointer-events: none;
            user-select: none;
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
            fill: rgb(51, 51, 51);
        }

        .collapse-button-disabled_${unique_id} {
            fill: #eee5;
            stroke: #3333334f;
            stroke-width: 1px;
            cursor: not-allowed;
        }

        .collapse-button-disabled_${unique_id}:hover {
            fill: #eee6;
            cursor: not-allowed;
        }

        .collapse-icon-disabled_${unique_id} {
            pointer-events: none;
            user-select: none;
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
            fill: #3333334f;
            cursor: not-allowed;
        }

        /* Edge hover styles */
        .edge-group_${unique_id} path {
            pointer-events: stroke;
        }

        .edge-label_${unique_id} {
            pointer-events: all;
        }

        @keyframes containerHighlight {
            0% { 
                filter: drop-shadow(0 0 0px rgba(255, 215, 0, 0));
            }
            25% { 
                filter: drop-shadow(0 0 15px rgba(255, 215, 0, 0.8));
            }
            50% { 
                filter: drop-shadow(0 0 0px rgba(255, 215, 0, 0));
            }
            75% { 
                filter: drop-shadow(0 0 15px rgba(255, 215, 0, 0.8));
            }
            100% { 
                filter: drop-shadow(0 0 0px rgba(255, 215, 0, 0));
            }
        }

        @keyframes containerStrokeHighlight {
            0% { 
                stroke-width: 1px;
            }
            25% { 
                stroke-width: 3px;
            }
            50% { 
                stroke-width: 1px;
            }
            75% { 
                stroke-width: 3px;
            }
            100% { 
                stroke-width: 1px;
            }
        }

        .highlight-cluster_${unique_id} > rect:first-of-type {
            animation: 
                containerStrokeHighlight 2.0s ease-in-out forwards;
        }

        .highlight-node_${unique_id} rect {
            animation: 
                containerStrokeHighlight 2.0s ease-in-out forwards;
        }

        text.node-label_${unique_id} {
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
            font-size: 0.8rem;  /* Base size using rem units for accessibility */
            cursor: default;
            user-select: none;
        }

        .node_${unique_id}.constant_${unique_id} .node-label_${unique_id} {
            font-size: 0.7rem;
        }

        .popup-content_${unique_id} summary {
            cursor: pointer;
            margin: 5px;
            user-select: none;
        }

        .node_${unique_id}.input_${unique_id} rect,
        .node_${unique_id}.output_${unique_id} rect,
        .node_${unique_id}.constant_${unique_id} rect {
            stroke-width: 1px;
            rx: 8px;
            ry: 8px;
            stroke-dasharray: none;
        }

        #png-error-container_${unique_id} {
            text-align: center;
            padding: 10px;
            border-bottom: 1px solid #ccc;
            font-size: 14px;
        }

        #png-error-message_${unique_id} {
            padding: 8px 16px;
            background-color: #dc3545;
            color: white;
            border-radius: 4px;
            display: inline-block;
            font-family: monospace;
            white-space: pre-wrap;
            text-align: left;
        }

        #png-link-container_${unique_id} {
            padding: 10px;
            border-bottom: 1px solid #ccc;
            font-size: 18px;
        }

        #png-download-button_${unique_id} {
            display: inline-block;
            padding: 8px 16px;
            background-color: #007bff;
            color: white;
            text-decoration: none;
            border-radius: 4px;
        }

        #png-download-button_${unique_id}:hover {
            background-color: #0056b3;
        }

        .custom-tooltip_${unique_id} {
            position: absolute;
            background: white;
            border: 1px solid #ccc;
            border-radius: 4px;
            padding: 6px 8px;
            font-size: 15px;
            font-family: inherit;
            box-shadow: 0 2px 8px rgba(0, 0, 0, 0.15);
            z-index: 1000;
            pointer-events: none;
            max-width: 350px;
            word-wrap: break-word;
            display: none;
        }
    </style>
    
    <script>
        var _module = window.module;
        var _exports = window.exports;
        var _define = window.define;
        
        window.module = undefined;
        window.exports = undefined;
        window.define = undefined;
        ${d3_source}
        window.module = _module;
        window.exports = _exports;
        window.define = _define;
    </script>
    <script>
        var _module = window.module;
        var _exports = window.exports;
        var _define = window.define;
        
        window.module = undefined;
        window.exports = undefined;
        window.define = undefined;
        
        ${viz_source}
        
        window.module = _module;
        window.exports = _exports;
        window.define = _define;
    </script>

    <style>
        ${jsoneditor_css}
    </style>

    <script>
        var _module = window.module;
        var _exports = window.exports;
        var _define = window.define;
        
        window.module = undefined;
        window.exports = undefined;
        window.define = undefined;
        
        ${jsoneditor_source}
        
        window.module = _module;
        window.exports = _exports;
        window.define = _define;
    </script>

    <script>


(function() {
    const adj_list = ${adj_list_json};
    const module_info = ${module_info_json};
    const func_info = ${func_info_json};
    const parent_module_to_nodes = ${parent_module_to_nodes_json};
    const parent_module_to_depth = ${parent_module_to_depth_json};
    const graph_node_name_to_without_suffix = ${graph_node_name_to_without_suffix};
    const ancestor_map = ${ancestor_map};
    const collapse_modules_after_depth = ${collapse_modules_after_depth};
    const node_to_module_path = ${node_to_module_path};
    const generateImage = ${generate_image};
    var collapsedContainers = new Set();
    const containerNodes = new Set();
    var adj_list_collapsed_nodes = {};
    let containerHierarchy = {};
    let vizInstance = null;
    let containerNodeCounts = {};

    const EXPAND_TARGET_COVERAGE = 0.8;
    const COLLAPSE_TARGET_COVERAGE = 0.2;
    let lastModifiedContainer = null; // Track the last modified container
    let lastOperation = null; // Track whether it was 'expand' or 'collapse'
    
    function waitForLibs(callback) {
        if (typeof d3 !== 'undefined' && typeof Viz !== 'undefined') {
            Viz.instance().then(viz => {
                vizInstance = viz;
                callback();
            });
        } else {
            setTimeout(function() { waitForLibs(callback) }, 400);
        }
    }

    let currentTooltip = null;

    function createTooltip(text) {
        removeTooltip();
        
        const tooltip = document.createElement('div');
        tooltip.className = `custom-tooltip_${unique_id}`;
        tooltip.textContent = text;
        
        const container = document.getElementById(`graph-container_${unique_id}`);
        container.appendChild(tooltip);
        currentTooltip = tooltip;
        
        return tooltip;
    }

    function positionTooltip(tooltip, buttonElement) {
        // Get the button's position in screen coordinates
        const buttonRect = buttonElement.getBoundingClientRect();
        
        // Get the container's position in screen coordinates
        const containerRect = document.getElementById(`graph-container_${unique_id}`).getBoundingClientRect();
        
        // Calculate position relative to container
        const buttonX = buttonRect.left + (buttonRect.width / 2) - containerRect.left;
        const buttonY = buttonRect.top - containerRect.top;
        
        // Show tooltip to measure its size
        tooltip.style.display = 'block';
        const tooltipRect = tooltip.getBoundingClientRect();
        
        // Position tooltip centered above button
        const left = buttonX - (tooltipRect.width / 2);
        const top = buttonY - tooltipRect.height - 6; // 6px gap above button
        
        tooltip.style.left = `${left}px`;
        tooltip.style.top = `${top}px`;
    }

    function showTooltip(buttonElement, text) {
        const tooltip = createTooltip(text);
        positionTooltip(tooltip, buttonElement);
    }

    function removeTooltip() {
        if (currentTooltip) {
            currentTooltip.remove();
            currentTooltip = null;
        }
    }

    function extractGraphData(graphData) {
        const result = {
            nodes: {},
            edges: [],
            clusters: {}
        };

        // Process nodes - they have _gvid property and may have a "pos" attribute
        for (const obj of graphData.objects) {
            // Skip cluster objects (which are subgraphs)
            if (obj.name && obj.name.startsWith("cluster_")) {
                result.clusters[obj.name.replace("cluster_", "")] = {
                    bb: obj.bb, // Bounding box "x1,y1,x2,y2"
                    label: obj.label,
                    nodes: obj.nodes ? obj.nodes.map(id => graphData.objects[id].name) : []
                };
                continue;
            }
        
            // This is a node
            if (obj.pos) {
            const [x, y] = obj.pos.split(",").map(Number);
                result.nodes[obj.name] = {
                id: obj._gvid,
                name: obj.name,
                position: { x, y },
                    width: parseFloat(obj.width) * 72, // Convert to pixels (72 points per inch)
                    height: parseFloat(obj.height) * 72,
                    label: obj.label === "\\N" ? obj.name : obj.label
                };
            }
        }

        // Process edges
        for (const edge of graphData.edges) {
            const sourceId = edge.tail;
            const targetId = edge.head;

            // Find node names by their _gvid
            const sourceName = Object.values(result.nodes).find(n => n.id === sourceId)?.name;
            const targetName = Object.values(result.nodes).find(n => n.id === targetId)?.name;

            if (!sourceName || !targetName) continue;

            // Extract path points from the pos attribute
            let pathPoints = [];
            if (edge.pos) {
                // IMPORTANT CHANGE: Correctly handle the "e," prefix
                const posStr = edge.pos;
                // Check if the string starts with "e,"
                if (posStr.startsWith("e,")) {
                    // Split the string and extract coordinates correctly
                    const parts = posStr.split(" ");
                    // Extract the endpoint (first coordinate after "e,")
                    const endPoint = parts[0].substring(2).split(",").map(Number);
                    // Add the remaining points
                    pathPoints = parts.slice(1).map(point => {
                        const [x, y] = point.split(',').map(Number);
                        return { x, y };
                    });

                    // Add the endpoint as the last point in the path
                    pathPoints.push({ x: endPoint[0], y: endPoint[1] });
                } else {
                    // Handle normal path without "e," prefix
                    pathPoints = posStr.split(' ').map(point => {
                        const [x, y] = point.split(',').map(Number);
                        return { x, y };
                    });
                }
            }
        
            let labelPos = null;
            if (edge.lp) {
                const [lx, ly] = edge.lp.split(',').map(Number);
                labelPos = { x: lx, y: ly };
            }

            result.edges.push({
                source: sourceName,
                target: targetName,
                path: pathPoints,
                edge_data_id: edge.edge_data_id,
                labelPos: labelPos
            });
        }

        return result;
    }

    function identifyContainersWithFailedNodes() {
        const failedNodes = Object.keys(adj_list).filter(node => adj_list[node].failed);
      
        const containersWithFailedNodes = new Set();
      
        failedNodes.forEach(failedNode => {
            let currentNode = failedNode;

            while (ancestor_map[currentNode]) {
                const container = ancestor_map[currentNode];
                containersWithFailedNodes.add(container);
                currentNode = container;
            }
        });
      
        return containersWithFailedNodes;
    }

    function processCollapsedGraph() {
        adj_list_collapsed_nodes = {};
        // Build container hierarchy
        containerHierarchy = {};
        for (const [node, parent] of Object.entries(ancestor_map)) {
            if (parent && !containerHierarchy[parent]) {
                containerHierarchy[parent] = [];
            }
            if (parent) {
                containerHierarchy[parent].push(node);
            }
        }

        // Helper function to get all descendants of a container
        function getAllDescendants(container) {
            let descendants = [];
            const children = containerHierarchy[container] || [];
            for (const child of children) {
                descendants.push(child);
                if (containerNodes.has(child)) {
                    descendants = descendants.concat(getAllDescendants(child));
                }
            }
            return descendants;
        }
        
        // First, add all collapsed containers as nodes
        for (const container of collapsedContainers) {
            adj_list_collapsed_nodes[container] = {
                node_type: "Module",
                edges: []
            };
        }
        
        for (const [node, data] of Object.entries(adj_list)) {
            let nodeIsInsideCollapsed = false;
            
            // Check if this node is inside a collapsed container
            for (const container of collapsedContainers) {
                const descendants = getAllDescendants(container);
                if (descendants.includes(node)) {
                    nodeIsInsideCollapsed = true;
                    break;
                }
            }
            
            // If node is not inside a collapsed container, add it to the processed list
            if (!nodeIsInsideCollapsed) {
                adj_list_collapsed_nodes[node] = {
                    node_type: data.node_type,
                    edges: []
                };
            }
        }

        function findVisibleCollapsedContainers() {
            const allDescendants = new Set();
            const visibleContainers = new Set();

            for (const container of collapsedContainers) {
                const descendants = getAllDescendants(container, containerHierarchy);
                for (const desc of descendants) {
                    allDescendants.add(desc);
                }
            }

            for (const container of collapsedContainers) {
                if (!allDescendants.has(container)) {
                    visibleContainers.add(container);
                }
            }

            return visibleContainers;
        }

        

        let visibleCollapsedContainers = findVisibleCollapsedContainers();
        
        // Process edges
        for (const [node, data] of Object.entries(adj_list)) {
            let seenEdgeDataIds = {};

            for (const edge of data.edges) {
                let sourceNode = node;
                let targetNode = edge.target;
                let skipEdge = false;

                for (const collapsedContainer of visibleCollapsedContainers) {
                    const descendants = getAllDescendants(collapsedContainer);

                    if (descendants.includes(node)) {
                        sourceNode = collapsedContainer;
                    }

                    if (descendants.includes(edge.target)) {
                        targetNode = collapsedContainer;
                    }

                    // Skip internal edges within collapsed containers
                    if (descendants.includes(node) && descendants.includes(edge.target)) {
                        skipEdge = true;
                    }
                }

                if (skipEdge || sourceNode === targetNode || !adj_list_collapsed_nodes[sourceNode]) {
                    continue;
                }

                // Create a consistent key for the source-target pair
                const pairKey = `${sourceNode}->${targetNode}`;
                const edgeId = edge.edge_data_id;

                if (edgeId) {
                    if (!seenEdgeDataIds[pairKey]) {
                        seenEdgeDataIds[pairKey] = new Set();
                    }

                    // Skip if this edge_data_id is already added between this pair
                    if (seenEdgeDataIds[pairKey].has(edgeId)) {
                        continue;
                    }

                    seenEdgeDataIds[pairKey].add(edgeId);
                }

                adj_list_collapsed_nodes[sourceNode].edges.push({
                    target: targetNode,
                    dims: edge.dims,
                    edge_data_id: edge.edge_data_id,
                });
            }
        }
    }

    function generateDotFromProcessedData(topDown = false) {
        function getNodeColor(nodeName) {
            if (adj_list_collapsed_nodes[nodeName]) {
                const nodeData = adj_list_collapsed_nodes[nodeName];
                if (nodeData.node_type === "Input") return "#FF9AFF";
                if (nodeData.node_type === "Constant") return "#EBEBEB";
                if (nodeData.node_type === "Output") return "#74e071";
                if (nodeData.node_type === "Module") return "#8BDEDB";
                if (nodeData.node_type === "Operation") return "#FEBC5A";
            }
            if (collapsedContainers.has(nodeName)) return "#8BDEDB";
            return "#FEBC5A";
        }

        function getErrorColor(nodeName) {
            if (adj_list[nodeName] && adj_list[nodeName].failed) return "#e00f0f";
            return null;
        }
        let dotSource = 'digraph G {\n';
        if (topDown) {
            dotSource += '  rankdir=TB;\n';
        } else {
            dotSource += '  rankdir=BT;\n';
        }
        dotSource += '  node [shape=box, style=filled];\n';
        dotSource += '  bgcolor="#f3f3eb";\n';
        dotSource += '  ranksep=0.5;\n';
        dotSource += '  nodesep=0.4;\n';


        // Find top containers
        const topContainers = new Set();
        for (const container of containerNodes) {
            if (!ancestor_map[container] & !collapsedContainers.has(container)) {
                topContainers.add(container);
            }
        }

        // Process containers recursively
        function processContainer(containerName, level = 0) {
            const indent = ' '.repeat(level * 2);
            
            // Check if this container is collapsed (appears as a regular node)
            if (collapsedContainers.has(containerName)) {
                // For collapsed containers, create a normal node instead of a subgraph
                const errorColor = getErrorColor(containerName);
                const fillColor = errorColor || getNodeColor(containerName);
                const displayName = graph_node_name_to_without_suffix[containerName] || containerName;
                dotSource += `${indent}"${containerName}" [fillcolor="${fillColor}", label="${displayName}"];\n`;
                return;
            }
            
            // Start subgraph for non-collapsed containers
            dotSource += `${indent}subgraph cluster_${containerName} {\n`;
            const displayName = graph_node_name_to_without_suffix[containerName] || containerName;
            dotSource += `${indent}  label="${displayName}";\n`;
            dotSource += `${indent}  fontsize=10;\n`;
            dotSource += `${indent}  style=rounded;\n`;
            dotSource += `${indent}  labelloc=t;\n`;
            dotSource += `${indent}  labeljust=l;\n`;

            // Get all direct children
            const children = containerHierarchy[containerName] || [];

            // Process child containers first
            for (const child of children) {
                if (containerNodes.has(child)) {
                    processContainer(child, level + 1);
                }
            }

            // Add regular nodes that are direct children
            for (const child of children) {
                if (!containerNodes.has(child) && adj_list_collapsed_nodes[child]) {
                    const nodeData = adj_list_collapsed_nodes[child];
                    if (nodeData.node_type === "Constant") {
                        dotSource += `${indent}  "${child}" [width=0.7, height=0.3];\n`;
                    } else if (nodeData.node_type === "Input" || nodeData.node_type === "Output") {
                        dotSource += `${indent}  "${child}" [width=1.2, height=0.5];\n`;
                    } else {
                        const errorColor = getErrorColor(child);
                        const fillColor = errorColor || getNodeColor(child);
                        const childDisplayName = graph_node_name_to_without_suffix[child] || child;
                        dotSource += `${indent}  "${child}" [fillcolor="${fillColor}", label="${childDisplayName}"];\n`;
                    }
                }
            }

            dotSource += `${indent}}\n`;
        }

        // Process all top containers
        for (const container of topContainers) {
            processContainer(container);
        }

        // Add standalone nodes (not in any container)
        for (const [node, data] of Object.entries(adj_list_collapsed_nodes)) {
            if (!ancestor_map[node]) {
                if (data.node_type === "Constant") {
                    const errorColor = getErrorColor(node);
                    const fillColor = errorColor || getNodeColor(node);
                    const nodeDisplayName = graph_node_name_to_without_suffix[node] || node;
                    dotSource += `  "${node}" [width=0.7, height=0.3, fillcolor="${fillColor}", label="${nodeDisplayName}"];\n`;
                } else if (data.node_type === "Input" || data.node_type === "Output") {
                    const errorColor = getErrorColor(node);
                    const fillColor = errorColor || getNodeColor(node);
                    const nodeDisplayName = graph_node_name_to_without_suffix[node] || node;
                    dotSource += `  "${node}" [width=1.2, height=0.5, fillcolor="${fillColor}", label="${nodeDisplayName}"];\n`;
                } else {
                    const errorColor = getErrorColor(node);
                    const fillColor = errorColor || getNodeColor(node);
                    const nodeDisplayName = graph_node_name_to_without_suffix[node] || node;
                    dotSource += `  "${node}" [fillcolor="${fillColor}", label="${nodeDisplayName}"];\n`;
                }
            }
        }

        // Add all edges
        for (const [node, data] of Object.entries(adj_list_collapsed_nodes)) {
            for (const edge of data.edges) {
                let edgeLabelFontSize = "10";
                if (edge.edge_data_id) {
                    dotSource += `  "${node}" -> "${edge.target}"[label="${edge.dims}", fontsize="${edgeLabelFontSize}", edge_data_id="${edge.edge_data_id}" ];\n`;
                } else {
                    dotSource += `  "${node}" -> "${edge.target}"[label="${edge.dims}", fontsize="${edgeLabelFontSize}" ];\n`;
                }
            }
        }

        dotSource += '}';
        return dotSource;
    }

    // Updated generateDotSource function
    function generateDotSource() {
        processCollapsedGraph();
        const dotSource = generateDotFromProcessedData(generateImage);
        return dotSource;
    }

    function calculateBoundingBox(nodeGroup, nodePositions) {
        let minX = Infinity, minY = Infinity, maxX = -Infinity, maxY = -Infinity;
        nodeGroup.forEach(nodeName => {
            const node = nodePositions[nodeName];
            if (node) {
                minX = Math.min(minX, node.x - node.width / 2);
                minY = Math.min(minY, node.y - node.height / 2);
                maxX = Math.max(maxX, node.x + node.width / 2);
                maxY = Math.max(maxY, node.y + node.height / 2);
            }
        });
        return { x: minX, y: minY, width: maxX - minX, height: maxY - minY };
    }

    function initContainerNodes() {
        // Clear and rebuild the containerNodes set
        containerNodes.clear();
        for (const value of Object.values(ancestor_map)) {
            if (value !== null && value !== undefined) {
                containerNodes.add(value);
            }
        }
    }

    function calculateContainerNodeCounts() {
        containerNodeCounts = {};
        
        for (const container of containerNodes) {
            containerNodeCounts[container] = 0;
        }
        
        for (const [node, parent] of Object.entries(ancestor_map)) {
            if (parent) {
                containerNodeCounts[parent] = (containerNodeCounts[parent] || 0) + 1;
                
                let currentParent = ancestor_map[parent];
                while (currentParent && !collapsedContainers.has(parent)) {
                    containerNodeCounts[currentParent] = (containerNodeCounts[currentParent] || 0) + 1;
                    currentParent = ancestor_map[currentParent];
                }
            }
        }
    }

    function reloadGraph() {
        d3.select("#graph_${unique_id}").html("");

        const dotSource = generateDotSource(adj_list, ancestor_map);

        graphData = vizInstance.renderJSON(dotSource)
        // Extract the graph data with positions from Graphviz
        const extractedGraphData = extractGraphData(graphData);

        if (!generateImage) {
            createGraph(extractedGraphData);

            // Needed to get the right vp coverage when containers are expanded
            calculateContainerNodeCounts();
        } else {
            renderGraphvizImage(dotSource);
        }
    }
                
    function initializeCollapsedContainersFromGraph() {
        const containerDepths = new Map();
        
        function getContainerDepth(container) {
            if (containerDepths.has(container)) {
                return containerDepths.get(container);
            }
            
            let depth = 0;
            for (const [child, parent] of Object.entries(ancestor_map)) {
                if (child === container && parent) {
                    depth = getContainerDepth(parent) + 1;
                    break;
                }
            }
            
            containerDepths.set(container, depth);
            return depth;
        }
        
        for (const container of containerNodes) {
            const depth = getContainerDepth(container);
            if (depth >= collapse_modules_after_depth) {
                collapsedContainers.add(container);
            }
        }
    }

    waitForLibs(function() {
        initContainerNodes();
        initializeCollapsedContainersFromGraph();
        reloadGraph();
    });

    function findErrorNode(nodes) {
        return nodes.find(node => adj_list[node].failed);
    }

    function renderGraphvizImage(dotSource) {
        d3.select("#png-link-container_${unique_id}").remove();
        
        function showErrorMessage(error) {
            const parent = d3.select("#graph-container_${unique_id}").node().parentNode;
            const errorContainer = d3.select(parent)
                .insert("div", "#graph-container_${unique_id}")
                .attr("id", "png-error-container_${unique_id}");
            
            errorContainer.append("div")
                .attr("id", "png-error-message_${unique_id}")
                .html(`<strong>Error:</strong><br>${error.stack || error.message || error}<br><br><strong>Note:</strong> Exporting very large models as PNG is currently not supported. Please use generate_image=False, or raise an issue on https://github.com/sachinhosmani/torchvista/issues`);
        }
        
        async function generatePNG() {
            processCollapsedGraph();
            const tbDotSource = generateDotFromProcessedData('TB');
            const svgString = vizInstance.renderString(tbDotSource, { format: 'svg' });
            
            const svgBlob = new Blob([svgString], { type: 'image/svg+xml' });
            const svgUrl = URL.createObjectURL(svgBlob);
            
            const img = new Image();
            const imageLoaded = new Promise((resolve, reject) => {
                img.onload = resolve;
                img.onerror = reject;
                img.src = svgUrl;
            });
            
            await imageLoaded;
            
            const canvas = document.createElement('canvas');
            const ctx = canvas.getContext('2d');
            const padding = 40;
            canvas.width = img.naturalWidth + (padding * 2);
            canvas.height = img.naturalHeight + (padding * 2);

            ctx.fillStyle = '#f3f3eb';
            ctx.fillRect(0, 0, canvas.width, canvas.height);
            ctx.drawImage(img, padding, padding);
            
            const pngBlob = await new Promise((resolve, reject) => {
                canvas.toBlob(resolve, 'image/png');
            });
            
            const pngUrl = URL.createObjectURL(pngBlob);
            
            // Create link container
            const parent = d3.select("#graph-container_${unique_id}").node().parentNode;
            const linkContainer = d3.select(parent)
                .insert("div", "#graph-container_${unique_id}")
                .attr("id", "png-link-container_${unique_id}");
            
            linkContainer.append("a")
                .attr("href", pngUrl)
                .attr("target", "_blank")
                .attr("id", "png-download-button_${unique_id}")
                .text("View PNG in New Tab");
                
            URL.revokeObjectURL(svgUrl);
        }
        
        generatePNG().catch(error => {
            console.error("Error generating PNG:", error);
            showErrorMessage(error);
        });
    }

    function createGraph(graphvizData) {
        const svg = d3.select("#graph_${unique_id}")
        .append("svg")
        .attr("width", "100%")
        .attr("height", "100%")
        .style("background-color", "#f3f3eb")
        .style("height", "100%");

        const inner = svg.append("g");

        // Add zoom behavior
        const zoom = d3.zoom().on("zoom", function(event) {
            removeTooltip();
            inner.attr("transform", event.transform);
            closePopup();
        });

        svg.call(zoom);

        // Define arrowhead marker
        svg.append("defs").append("marker")
        .attr("id", "arrowhead_${unique_id}")
        .attr("viewBox", "0 -5 10 10")
        .attr("refX", 8)
        .attr("refY", 0)
        .attr("markerWidth", 6)
        .attr("markerHeight", 6)
        .attr("orient", "auto")
        .append("path")
        .attr("d", "M0,-5L10,0L0,5")
        .attr("fill", "#333");

        // Draw clusters (container boxes) first so they're behind nodes
        const clusters = inner.selectAll(".boundary-box_${unique_id}")
        .data(Object.entries(graphvizData.clusters))
        .enter()
        .append("g")
        .attr("class", "boundary-box_${unique_id}");

        const containersWithFailedNodes = identifyContainersWithFailedNodes(adj_list, ancestor_map);

        clusters.each(function([clusterName, clusterData]) {
            const group = d3.select(this);

            const hasFailedNodes = containersWithFailedNodes.has(clusterName);
            if (hasFailedNodes) {
                group.classed("error_${unique_id}", true);
            }

            // Parse bounding box coordinates
            const [x1, y1, x2, y2] = clusterData.bb.split(",").map(Number);
            
            // Create boundary box
            group.append("rect")
            .attr("x", x1)
            .attr("y", y1)
            .attr("width", x2 - x1)
            .attr("height", y2 - y1)
            .attr("fill-opacity", 0.0)
            .attr("stroke", "black")
            .attr("stroke-width", 1)
            .attr("rx", 5)
            .attr("ry", 5);
                
            // Create label for the cluster
            const textGroup = group.append("g")
            .attr("transform", `translate(${x1 + 10}, ${y1 + 3})`);
            
            const text = textGroup.append("text")
            .text(graph_node_name_to_without_suffix[clusterData.label] || clusterData.label)
            .attr("text-anchor", "start")
            .attr("dominant-baseline", "hanging")
            .attr("dy", "-6")
            .attr("font-size", 7)
            .attr("fill", "#333")
            .attr("stroke", "#f3f3eb")
            .attr("stroke-width", "5px")
            .attr("paint-order", "stroke")
            .attr("stroke-linejoin", "round");
                        
            // Add collapse/expand button
            const collapseButton = group.append("circle")
            .attr("class", "collapse-button_${unique_id}")
            .attr("cx", x1 + 10)
            .attr("cy", y1 + 10)
            .attr("r", 5.5);

            group.append("text")
            .attr("class", "collapse-icon_${unique_id}")
            .attr("x", x1 + 10)
            .attr("y", y1 + 10)
            .attr("text-anchor", "middle")
            .attr("dominant-baseline", "middle")
            .attr("font-size", "10px")
            .text("-");

            // Add custom tooltip events
            collapseButton.on("mouseenter", function() {
                showTooltip(this, "Collapse module");
            })
            .on("mouseleave", removeTooltip);
        });

        // Add click event handler for collapse buttons
        clusters.selectAll(".collapse-button_${unique_id}, .collapse-icon_${unique_id}")
        .on("click", function(event, [clusterName, clusterData]) {
            event.stopPropagation(); // Prevent event from bubbling to SVG
            removeTooltip();
            // Collapse the container
            collapsedContainers.add(clusterName);
            
            // Track this container and operation for centering after reload
            lastModifiedContainer = clusterName;
            lastOperation = "collapse";
            reloadGraph();
        });
        
        // Create edges using the path points from Graphviz
        const edges = inner.selectAll(".edge_${unique_id}")
        .data(graphvizData.edges)
        .enter()
        .append("g")
        .attr("class", "edge-group_${unique_id}");
        
        // Add the actual path
        edges.append("path")
        .attr("class", "edge_${unique_id}")
        .attr("d", function(edge) {
            if (edge.path && edge.path.length > 0) {
                    // Create a proper SVG path using Bézier curves
                let d = "";
                const points = edge.path;
                
                    // Start at the first point
                d += `M${points[0].x},${points[0].y}`;
                
                    // Create cubic Bézier curves
                for (let i = 1; i < points.length - 2; i += 3) {
                    if (i + 2 < points.length) {
                        d += ` C${points[i].x},${points[i].y} ${points[i+1].x},${points[i+1].y} ${points[i+2].x},${points[i+2].y}`;
                    }
                }
                
                    // If we have leftover points, draw a line to the last point
                if (points.length % 3 !== 1) {
                    d += ` L${points[points.length-1].x},${points[points.length-1].y}`;
                }
                
                return d;
            }
            return "";
        })
        .attr("marker-end", "url(#arrowhead_${unique_id})");
        
        // Add dimension labels on edges
        edges.each(function(edge) {
            const edgeGroup = d3.select(this);
            const path = edgeGroup.select("path").node();
            
            if (!path) return;
            
            // Find the source node data
            if (!adj_list[edge.source] && !adj_list_collapsed_nodes[edge.source]) return;
            
            // Find the edge data with matching target
            let sourceNodeData = adj_list_collapsed_nodes[edge.source];
            if (!sourceNodeData) {
                sourceNodeData = adj_list[edge.source];
            }
            const edgeData = sourceNodeData.edges.find(e => 
                e.target === edge.target && 
                (e.edge_data_id === undefined || e.edge_data_id.toString() === edge.edge_data_id)
            );
            if (!edgeData || !edgeData.dims) return;
            
            const pointAtStart = edge.labelPos;

            const angle = 0.0;
            
            // Create a background for the text
            const labelGroup = edgeGroup.append("g")
            .attr("class", "edge-label_${unique_id}")
            .attr("transform", `translate(${pointAtStart.x}, ${pointAtStart.y})`);
            
            const text = labelGroup.append("text")
            .attr("text-anchor", "middle")
            .attr("dominant-baseline", "central")
            .style("font-size", "6pt")
            .text(edgeData.dims)
            // .attr("transform", `rotate(${angle})`)
            .attr("stroke", "#f3f3eb")
            .attr("stroke-width", "3pt")
            .attr("paint-order", "stroke")
            .attr("stroke-linejoin", "round");

            // Add background rectangle
            const bbox = text.node().getBBox();
            labelGroup.insert("rect", "text")
            .attr("x", bbox.x - 2)
            .attr("y", bbox.y - 2)
            .attr("rx", 4)
            .attr("ry", 4)
            .attr("width", bbox.width + 4)
            .attr("height", bbox.height + 4)
            .attr("fill", "rgb(243, 243, 235) ")
            .attr("stroke", "#79797933")
            .attr("stroke-width", "0.5px")
            .attr("fill-opacity", 1.0);
            // .attr("transform", `rotate(${angle})`);
        });
        

        // Add hover effects to edge paths
        edges.each(function() {
            const edgeGroup = d3.select(this);
            const path = edgeGroup.select("path");
            const labelGroup = edgeGroup.select(".edge-label_${unique_id}");
            const labelText = labelGroup.select("text");
            const labelBg = labelGroup.select("rect");
            const originalStrokeWidth = path.attr("stroke-width") || 1;
            const originalFontSize = labelText.style("font-size") || "6pt";
            const originalParent = edgeGroup.node().parentNode;
            const originalPosition = Array.from(originalParent.childNodes).indexOf(edgeGroup.node());
            
            function highlightEdge() {
                // Move to end of parent's children list to appear on top
                originalParent.appendChild(edgeGroup.node());
                
                // Make edge thicker
                path.attr("stroke-width", parseFloat(originalStrokeWidth) * 1.4);
                
                // Make label text larger
                labelText.style("font-size", "7pt");
                    
                // Adjust the background rectangle for the label
                const bbox = labelText.node().getBBox();
                labelBg.attr("x", bbox.x - 2)
                    .attr("y", bbox.y - 2)
                    .attr("width", bbox.width + 4)
                    .attr("height", bbox.height + 4);
            }
            
            function resetEdge() {
                // Restore original position in parent's children list
                const currentIndex = Array.from(originalParent.childNodes).indexOf(edgeGroup.node());
                if (currentIndex !== originalPosition) {
                    // Insert back at original position
                    const referenceNode = originalParent.childNodes[originalPosition] || null;
                    originalParent.insertBefore(edgeGroup.node(), referenceNode);
                }
                
                // Restore original edge thickness
                path.attr("stroke-width", originalStrokeWidth);
                
                // Restore original label text size
                labelText.style("font-size", originalFontSize);
                    
                // Restore original background rectangle size
                const bbox = labelText.node().getBBox();
                labelBg.attr("x", bbox.x - 2)
                    .attr("y", bbox.y - 2)
                    .attr("width", bbox.width + 4)
                    .attr("height", bbox.height + 4);
            }
            
            // Add hover events to the path
            path.on("mouseenter", highlightEdge)
                .on("mouseleave", resetEdge);
            
            // Add hover events to label group (both text and background)
            if (labelGroup.node()) {
                labelText.on("mouseenter", highlightEdge)
                    .on("mouseleave", resetEdge);
                    
                labelBg.on("mouseenter", highlightEdge)
                    .on("mouseleave", resetEdge);
            }
        });


        // Create nodes using positions from Graphviz
        const nodes = inner.selectAll(".node_${unique_id}")
            .data(Object.entries(graphvizData.nodes))
            .enter()
            .append("g")
            // In the createGraph function, when you're creating nodes, modify the class assignment:
            .attr("class", function([nodeName, nodeData]) {
                let classes = "node_${unique_id}";
                if ((adj_list[nodeName] && adj_list[nodeName].failed) || containersWithFailedNodes.has(nodeName)) classes += " error_${unique_id}";
                if ((adj_list[nodeName] && adj_list[nodeName].node_type === "Input")) classes += " input_${unique_id}";
                else if ((adj_list[nodeName] && adj_list[nodeName].node_type === "Module") || (collapsedContainers.has(nodeName))) classes += " module_${unique_id}";
                else if ((adj_list[nodeName] && adj_list[nodeName].node_type === "Output")) classes += " output_${unique_id}";
                else if ((adj_list[nodeName] && adj_list[nodeName].node_type === "Operation")) classes += " operation_${unique_id}";
                else if ((adj_list[nodeName] && adj_list[nodeName].node_type === "Constant")) classes += " constant_${unique_id}";
                return classes;
            })
            .attr("transform", function([nodeName, nodeData]) {
                return `translate(${nodeData.position.x}, ${nodeData.position.y})`;
            });
        
        
        // Add rectangles for nodes
        nodes.append("rect")
        .attr("width", function([nodeName, nodeData]) { return nodeData.width; })
        .attr("height", function([nodeName, nodeData]) { return nodeData.height; })
        .attr("x", function([nodeName, nodeData]) { return -nodeData.width / 2; })
        .attr("y", function([nodeName, nodeData]) { return -nodeData.height / 2; })
        .attr("class", function([nodeName]) { 
            if (adj_list[nodeName] && adj_list[nodeName].node_type === "Input") return "input_${unique_id}";
            if (adj_list[nodeName] && adj_list[nodeName].node_type === "Output") return "output_${unique_id}";
            if (adj_list[nodeName] && adj_list[nodeName].node_type === "Constant") return "constant_${unique_id}";
            if ((adj_list[nodeName] && adj_list[nodeName].failed) || containersWithFailedNodes.has(nodeName)) return "failed_${unique_id}";
            return ((adj_list[nodeName] && adj_list[nodeName].node_type === "Module") || (collapsedContainers.has(nodeName))) ? "module_${unique_id}" : "operation_${unique_id}";
        });
        
        // Add labels for nodes
        nodes.append("text")
        .attr("class", "node-label_${unique_id}")
        .text(function([nodeName]) { return graph_node_name_to_without_suffix[nodeName] || nodeName; })
        .attr("text-anchor", "middle")
        .attr("dominant-baseline", "central");
        
        // Add type indicators
        nodes.filter(function([nodeName]) {
            return !(((adj_list[nodeName] && adj_list[nodeName].node_type === "Output")) || ((adj_list[nodeName] && adj_list[nodeName].node_type === "Input")) || ((adj_list[nodeName] && adj_list[nodeName].node_type === "Constant")));
        }).append("text")
        .attr("class", "type-indicator_${unique_id}")
        .attr("text-anchor", "start")
        .attr("dominant-baseline", "text-bottom")
        .attr("transform", function([nodeName, nodeData]) {
            return `translate(${-nodeData.width / 2 + 3}, ${nodeData.height / 2 - 3})`;
        })
        .text(function([nodeName]) {
            return (adj_list[nodeName] && adj_list[nodeName].node_type === "Module") || (collapsedContainers.has(nodeName)) ? "Module" : "Tensor Op";
        });

        // Add info buttons to nodes (except input, output, constant)
        inner.selectAll(".node_${unique_id}")
        .filter(function([nodeName]) {
            return !(((adj_list[nodeName] && adj_list[nodeName].node_type === "Output")) || ((adj_list[nodeName] && adj_list[nodeName].node_type === "Input")) || ((adj_list[nodeName] && adj_list[nodeName].node_type === "Constant")))
        })
        .each(function([nodeName, nodeData]) {
            const nodeGroup = d3.select(this);
            const width = nodeData.width;
            const height = nodeData.height;
            
        // Add info button circle
        const infoButton = nodeGroup.append("circle")
        .attr("class", "info-button_${unique_id}")
        .attr("cx", width/2 - 7)
        .attr("cy", -height/2 + 7)
        .attr("r", 5.5);

        // Add info button icon
        nodeGroup.append("text")
        .attr("class", "info-icon_${unique_id}")
        .attr("x", width/2 - 7)
        .attr("y", -height/2 + 8)
        .attr("text-anchor", "middle")
        .attr("dominant-baseline", "middle")
        .attr("font-size", "8px")
        .attr("font-weight", "bold")
        .text("i");

        // Add custom tooltip
        infoButton.on("mouseenter", function() {
            showTooltip(this, "Show info");
        })
        .on("mouseleave", removeTooltip);

        // Add click handler to info button
            nodeGroup.selectAll(".info-button_${unique_id}, .info-icon_${unique_id}")
            .on("click", function(event) {
                event.stopPropagation();
                const infoButton = nodeGroup.select(".info-button_${unique_id}").node();
                showPopupFromInfoButton.call(nodeGroup.node(), event, [nodeName], infoButton);
            });
        });

        // Add expand buttons to all module nodes (including non-expandable ones)
        inner.selectAll(".node_${unique_id}")
        .filter(function([nodeName]) {
            return (adj_list[nodeName] && adj_list[nodeName].node_type === "Module") || collapsedContainers.has(nodeName);
        })
        .each(function([nodeName, nodeData]) {
            const nodeGroup = d3.select(this);
            const width = nodeData.width;
            const height = nodeData.height;
            
            // Check if this module is expandable (has children)
            const isExpandable = collapsedContainers.has(nodeName);
            
            // Calculate depth of this node
            let depth = 0;
            let currentNode = nodeName;
            while (ancestor_map[currentNode]) {
                depth++;
                currentNode = ancestor_map[currentNode];
            }
            
            // Add expand button circle
            const expandButton = nodeGroup.append("circle")
            .attr("class", isExpandable ? "collapse-button_${unique_id}" : "collapse-button-disabled_${unique_id}")
            .attr("cx", -width / 2 + 7)
            .attr("cy", -height / 2 + 7)
            .attr("r", 5.5)
            .attr("stroke-width", 1);


            // Add expand button icon
            nodeGroup.append("text")
            .attr("class", isExpandable ? "collapse-icon_${unique_id}" : "collapse-icon-disabled_${unique_id}")
            .attr("x", -width / 2 + 7)
            .attr("y", -height / 2 + 7)
            .attr("text-anchor", "middle")
            .attr("dominant-baseline", "middle")
            .attr("font-size", "10px")
            .text("+");

            // Add custom tooltip
            const tooltipText = isExpandable ? 
                "Expand module" : 
                `Call with forced_module_tracing_depth=${depth + 1} or more to make this module expandable`;

            expandButton.on("mouseenter", function() {
                showTooltip(this, tooltipText);
            })
            .on("mouseleave", removeTooltip);

            // Only add click handler for expandable modules
            if (isExpandable) {
                nodeGroup.selectAll(".collapse-button_${unique_id}, .collapse-icon_${unique_id}")
                .on("click", function(event) {
                    event.stopPropagation();
                    removeTooltip();
                    
                    // Expand the module
                    collapsedContainers.delete(nodeName);

                    // Track this container and operation for centering after reload
                    lastModifiedContainer = nodeName;
                    lastOperation = "expand";
                    
                    // Use the existing reloadGraph function
                    reloadGraph();
                });
            }
        });
        
        // Add hover effects
        nodes.each(function() {
            const nodeGroup = d3.select(this);
            const rect = nodeGroup.select("rect");
            const originalColor = rect.style("fill");
            let isHovered = false;
            
            nodeGroup.on("mouseenter", function() {
                if (!isHovered) {
                    isHovered = true;
                    const currentColor = d3.color(rect.style("fill") || originalColor);
                    rect.style("fill", currentColor.brighter(0.4));
                }
            })
            .on("mouseleave", function() {
                isHovered = false;
                rect.style("fill", originalColor);
            });
        });

        function createJsonViewer(container, data) {
            // Clear previous contents
            container.innerHTML = '';

            // Create a new div to hold the editor
            const editorContainer = document.createElement('div');
            editorContainer.style.height = '180px';
            container.appendChild(editorContainer);

            const options = {
                mode: 'code',
                mainMenuBar: false,
                navigationBar: false,
                statusBar: false,
                indentation: 2,
                onChange: function() {
                    editor.set(data);
                }
            };
            const editor = new JSONEditor(editorContainer, options);
            editor.set(data);
            // set a reference for later use during resizing
            container.jsonEditor = editor;
        }
        
        function showPopupFromInfoButton(event, [nodeName], infoButtonElement = null) {
            event.stopPropagation();
            
            const clickedNodeElement = d3.select(this);

            // Position based on info button or node center
            const clickednodeData = graphvizData.nodes[nodeName];
            if (clickednodeData) {
                const svgNode = svg.node();
                const width = svgNode.clientWidth || 800;
                const height = svgNode.clientHeight || 600;
                
                const currentTransform = d3.zoomTransform(svg.node());
                
                let targetX, targetY;
                
                if (infoButtonElement) {
                    // Calculate info button position relative to node center
                    const infoButtonOffsetX = clickednodeData.width / 2 - 8;
                    const infoButtonOffsetY = -clickednodeData.height / 2 + 8;
                    
                    // Calculate where the info button should be positioned on screen
                    const infoButtonWorldX = clickednodeData.position.x + infoButtonOffsetX;
                    const infoButtonWorldY = clickednodeData.position.y + infoButtonOffsetY;
                    
                    targetX = 0.4 * width - infoButtonWorldX * currentTransform.k;
                    targetY = 0.2 * height - infoButtonWorldY * currentTransform.k;
                } else {
                    // Original behavior - center the node
                    targetX = 0.4 * width - clickednodeData.position.x * currentTransform.k;
                    targetY = 0.2 * height - clickednodeData.position.y * currentTransform.k;
                }
                
                svg.call(zoom.transform, d3.zoomIdentity
                    .translate(targetX, targetY)
                    .scale(currentTransform.k));
            }

            clickedNodeElement.classed("highlight-node_${unique_id}", true);
            
            setTimeout(() => {
                clickedNodeElement.classed("highlight-node_${unique_id}", false);
            }, 2000);

            const popup = d3.select("#popup_${unique_id}");
            const moduleTemplate = d3.select("#module-template_${unique_id}");
            const tensorOpTemplate = d3.select("#tensor-op-template_${unique_id}");
            const noInfoMessage = d3.select("#no-info_${unique_id}");
            
            const nodeData = adj_list[nodeName];
            
            let popupTitle = nodeName;
            if (node_to_module_path[nodeName] && graph_node_name_to_without_suffix[nodeName]) {
                popupTitle = node_to_module_path[nodeName] + "." + graph_node_name_to_without_suffix[nodeName];
            }
            d3.select("#popup-title_${unique_id}").text(popupTitle);

            moduleTemplate.style("display", "none");
            tensorOpTemplate.style("display", "none");
            noInfoMessage.style("display", "none");
            
            if (module_info[nodeName]) {
                const info = module_info[nodeName];
                moduleTemplate.style("display", "block").classed("content-shown_${unique_id}", true);
                
                d3.select("#module-extra-repr_${unique_id}").text(info.extra_repr || 'N/A');

                const forward_info = func_info[nodeName];
                const attributesContainer = d3.select("#module-attributes_${unique_id}").node();
                attributesContainer.innerHTML = '';
                createJsonViewer(attributesContainer, info.attributes);

                const parametersContainer = d3.select("#module-parameters_${unique_id}").node();
                parametersContainer.innerHTML = '';
                createJsonViewer(parametersContainer, forward_info);
            } else if (!(adj_list[nodeName] && adj_list[nodeName].node_type === "Input") && func_info[nodeName]) {
                const info = func_info[nodeName];
                tensorOpTemplate.style("display", "block").classed("content-shown_${unique_id}", true);
                const parametersContainer = d3.select("#tensor-op-parameters_${unique_id}").node();
                parametersContainer.innerHTML = '';
                createJsonViewer(parametersContainer, info);
            } else {
                noInfoMessage.style("display", "block").classed("content-shown_${unique_id}", true);
            }
            
            // Position popup based on info button or node center
            let boundingBox, x, y;
            const containerRect = d3.select("#graph-container_${unique_id}").node().getBoundingClientRect();
            
            if (infoButtonElement) {
                boundingBox = infoButtonElement.getBoundingClientRect();
                x = boundingBox.left + boundingBox.width / 2 - containerRect.left;
                y = boundingBox.top + boundingBox.height / 2 - containerRect.top;
            } else {
                const clickedElement = d3.select(this).select("rect").node();
                boundingBox = clickedElement.getBoundingClientRect();
                x = boundingBox.left + boundingBox.width / 2 - containerRect.left;
                y = boundingBox.top + boundingBox.height / 2 - containerRect.top;
            }
            
            popup.style("left", `${x}px`)
            .style("top", `${y}px`)
            .style("display", "flex");
            d3.select("#popup-content_${unique_id}").node().scrollTop = 0;
        }

        function closePopup() {
            d3.select("#popup_${unique_id}").style("display", "none");
        }
        
        d3.select(".popup-close_${unique_id}").on("click", closePopup);
        svg.on("click", closePopup);

        const popup = document.getElementById(`popup_${unique_id}`);
        const resizeObserver = new ResizeObserver(() => {
            const editorContainers = popup.querySelectorAll('.json-viewer_${unique_id}');
            for (let container of editorContainers) {
                if (container.jsonEditor) {
                    container.jsonEditor.resize();
                }
            }
        });
        resizeObserver.observe(popup);
        
        // Calculate view box to ensure the graph is properly scaled and centered
        const findErrorNode = (nodes) => {
            for (const [nodeName] of nodes) {
                if (adj_list[nodeName] && adj_list[nodeName].failed) {
                    return nodeName;
                }
            }
            return null;
        };
        
        const errorNode = findErrorNode(Object.entries(graphvizData.nodes));
        
        // Initial zoom setup
        const svgNode = svg.node();
        const width = svgNode.clientWidth || 800; // Fallback if clientWidth is 0
        const height = svgNode.clientHeight || 600; // Fallback if clientHeight is 0
        
        if (lastModifiedContainer && (Object.keys(graphvizData.nodes).includes(lastModifiedContainer) || 
            Object.keys(graphvizData.clusters).includes(lastModifiedContainer))) {
            
            // Center on the container that was just expanded or collapsed
            let targetCoverage;
            if (lastOperation === "expand") {
                const nodeCount = containerNodeCounts[lastModifiedContainer] || 1;

                const minCoverage = 0.3;
                const maxNodes = 20;
                const scaleFactor = Math.min(nodeCount, maxNodes) / maxNodes;
                targetCoverage = Math.min(EXPAND_TARGET_COVERAGE, minCoverage + (EXPAND_TARGET_COVERAGE - minCoverage) * scaleFactor);
            } else {
                targetCoverage = COLLAPSE_TARGET_COVERAGE;
            }
            let targetNode;
            let targetPos;
            let containerWidth, containerHeight;
            
            if (Object.keys(graphvizData.nodes).includes(lastModifiedContainer)) {
                // It's a node (collapsed container)
                targetNode = graphvizData.nodes[lastModifiedContainer];
                targetPos = { x: targetNode.position.x, y: targetNode.position.y };
                containerWidth = targetNode.width;
                containerHeight = targetNode.height;
            } else if (Object.keys(graphvizData.clusters).includes(lastModifiedContainer)) {
                // It's a cluster (expanded container)
                const clusterData = graphvizData.clusters[lastModifiedContainer];
                const [x1, y1, x2, y2] = clusterData.bb.split(",").map(Number);
                targetPos = { x: (x1 + x2) / 2, y: (y1 + y2) / 2 };
                containerWidth = x2 - x1;
                containerHeight = y2 - y1;
            }
            
            if (targetPos && containerWidth && containerHeight) {
                // Calculate scale to make container occupy the target percentage of viewport
                const widthScale = (width * targetCoverage) / containerWidth;
                const heightScale = (height * targetCoverage) / containerHeight;
                // Use the smaller scale to ensure container fits in viewport
                const targetScale = Math.min(widthScale, heightScale);
                
                const translate = [
                    width / 2 - targetPos.x * targetScale,
                    height / 2 - targetPos.y * targetScale
                ];
                
                svg.call(zoom.transform, d3.zoomIdentity
                    .translate(translate[0], translate[1])
                    .scale(targetScale));
                
                // Apply highlight effect to the target container
                if (Object.keys(graphvizData.nodes).includes(lastModifiedContainer)) {
                    // Highlight the node (collapsed container)
                    const targetElement = inner.selectAll(".node_${unique_id}")
                        .filter(([nodeName]) => nodeName === lastModifiedContainer);
                        
                    if (!targetElement.empty()) {
                        targetElement.classed("highlight-node_${unique_id}", true);
                        
                        // Remove the highlight class after animation completes
                        setTimeout(() => {
                            targetElement.classed("highlight-node_${unique_id}", false);
                        }, 2000); // Match the animation duration
                    }
                } else if (Object.keys(graphvizData.clusters).includes(lastModifiedContainer)) {
                    // Highlight the cluster (expanded container)
                    const targetElement = inner.selectAll(".boundary-box_${unique_id}")
                        .filter(function([clusterName]) {
                            return clusterName === lastModifiedContainer;
                        });
                        
                    if (!targetElement.empty()) {
                        targetElement.classed("highlight-cluster_${unique_id}", true);
                        
                        // Remove the highlight class after animation completes
                        setTimeout(() => {
                            targetElement.classed("highlight-cluster_${unique_id}", false);
                        }, 2000); // Match the animation duration
                    }
                }
                
                // Reset tracking variables after applying the transformation
                lastModifiedContainer = null;
                lastOperation = null;
            }
        } else if (errorNode) {
            // Focus on the error node if there is one
            const nodeData = graphvizData.nodes[errorNode];
            if (nodeData) {
                const scale = 2; // Zoom in on error node
                const translate = [
                    width / 2 - nodeData.position.x * scale,
                    height / 2 - nodeData.position.y * scale
                    ];
                
                svg.call(zoom.transform, d3.zoomIdentity
                    .translate(translate[0], translate[1])
                    .scale(scale));
            }
        } else {
            // Show the entire graph
            // Find the bounds of the graph
            let minX = Infinity, minY = Infinity;
            let maxX = -Infinity, maxY = -Infinity;
            
            Object.values(graphvizData.nodes).forEach(node => {
                minX = Math.min(minX, node.position.x - node.width / 2);
                minY = Math.min(minY, node.position.y - node.height / 2);
                maxX = Math.max(maxX, node.position.x + node.width / 2);
                maxY = Math.max(maxY, node.position.y + node.height / 2);
            });
            
            // Also consider cluster bounding boxes
            Object.values(graphvizData.clusters).forEach(cluster => {
                const [x1, y1, x2, y2] = cluster.bb.split(",").map(Number);
                minX = Math.min(minX, x1);
                minY = Math.min(minY, y1);
                maxX = Math.max(maxX, x2);
                maxY = Math.max(maxY, y2);
            });
            
            const graphWidth = maxX - minX;
            const graphHeight = maxY - minY;
            
            // Determine scale to fit the entire graph with some padding
            const scale = Math.min(width / graphWidth, height / graphHeight) * 0.9;
            
            // Center the graph
            const translate = [
                width / 2 - (minX + graphWidth / 2) * scale,
                height / 2 - (minY + graphHeight / 2) * scale
                ];
            
            svg.call(zoom.transform, d3.zoomIdentity
                .translate(translate[0], translate[1])
                .scale(scale));
        }
            
    }
})();


    </script>
</div>