frontend/js/treemap.js

/**
 * LLM Scope - Pipeline Visualization
 * Professional, minimalist architecture diagrams
 */

class ModelTreeViz {
    constructor(containerId) {
        this.containerId = containerId;
        this.container = document.getElementById(containerId);
        this.pipelineData = null;
        this.totalParams = 0;
        this.globalMaxParams = null;

        // Layout
        this.baseNodeHeight = 48;
        this.minNodeHeight = 36;
        this.maxNodeHeight = 72;
        this.nodeGap = 12;
        this.containerPadding = 14;
        this.containerGap = 8;
        this.padding = 24;
        this.branchGap = 24;

        // Width
        this.minNodeWidth = 110;
        this.maxNodeWidth = 240;

        this.transform = d3.zoomIdentity;

        // Professional muted color palette
        this.colors = {
            embedding: '#d97706',   // amber
            attention: '#7c3aed',   // violet
            mlp: '#059669',         // emerald
            norm: '#0891b2',        // cyan
            head: '#db2777',        // pink
            layers: '#4f46e5',      // indigo
            encoder: '#6366f1',     // indigo lighter
            linear: '#6b7280',      // gray
            module: '#52525b',      // zinc
            model: '#3f3f46',       // zinc darker
            block: '#52525b',
            pooler: '#0d9488',      // teal
            parallel: '#3f3f46',
            default: '#52525b'
        };

        this._init();
    }

    _init() {
        this.svg = d3.select(`#${this.containerId}`)
            .append('svg')
            .attr('width', '100%')
            .attr('height', '100%');

        this.zoom = d3.zoom()
            .scaleExtent([0.2, 3])
            .on('zoom', (event) => {
                this.transform = event.transform;
                this.g.attr('transform', event.transform);
            });

        this.svg.call(this.zoom);

        this.g = this.svg.append('g');
        this.containersGroup = this.g.append('g').attr('class', 'containers');
        this.arrowsGroup = this.g.append('g').attr('class', 'arrows');
        this.nodesGroup = this.g.append('g').attr('class', 'nodes');

        // Arrow marker
        this.svg.append('defs').append('marker')
            .attr('id', `arrow-${this.containerId}`)
            .attr('viewBox', '0 -4 8 8')
            .attr('refX', 7)
            .attr('refY', 0)
            .attr('markerWidth', 5)
            .attr('markerHeight', 5)
            .attr('orient', 'auto')
            .append('path')
            .attr('d', 'M0,-4L8,0L0,4')
            .attr('fill', '#52525b');

        this.resizeObserver = new ResizeObserver(() => this._onResize());
        this.resizeObserver.observe(this.container);
    }

    _onResize() {
        if (this.pipelineData) {
            this._render();
        }
    }

    getColor(type) {
        return this.colors[type] || this.colors.default;
    }

    _getNodeWidth(params) {
        const refParams = this.globalMaxParams || this.totalParams;
        if (refParams === 0 || params === 0) return this.minNodeWidth;
        const ratio = params / refParams;
        const scale = Math.sqrt(ratio);
        return this.minNodeWidth + (this.maxNodeWidth - this.minNodeWidth) * scale;
    }

    _getNodeHeight(params) {
        const refParams = this.globalMaxParams || this.totalParams;
        if (refParams === 0 || params === 0) return this.minNodeHeight;
        const ratio = params / refParams;
        const scale = Math.pow(ratio, 0.25);
        return this.minNodeHeight + (this.maxNodeHeight - this.minNodeHeight) * scale;
    }

    setData(data, globalMaxParams = null) {
        this.pipelineData = data;
        this.totalParams = data.params || 0;
        this.globalMaxParams = globalMaxParams;
        this._render();
    }

    _calcStepWidth(step) {
        const nodeWidth = this._getNodeWidth(step.params || 0);

        if (step.type === 'parallel' && step.branches) {
            let totalWidth = 0;
            for (const branch of step.branches) {
                totalWidth += this._calcStepWidth(branch);
            }
            totalWidth += (step.branches.length - 1) * this.branchGap;
            return Math.max(nodeWidth, totalWidth);
        }

        if (step.substeps && step._collapsed === false) {
            let maxChildWidth = 0;
            for (const sub of step.substeps) {
                maxChildWidth = Math.max(maxChildWidth, this._calcStepWidth(sub));
            }
            return Math.max(nodeWidth, maxChildWidth + this.containerPadding * 2);
        }

        return nodeWidth;
    }

    _layoutPipeline(data) {
        const nodes = [];
        const arrows = [];
        const containers = [];
        const centerX = 350;

        const layoutSteps = (steps, startY, parentCenterX) => {
            if (!steps || steps.length === 0) {
                return { endY: startY, lastNodes: [] };
            }

            let y = startY;
            let prevNodes = [];

            for (let i = 0; i < steps.length; i++) {
                const step = steps[i];

                if (step.type === 'parallel' && step.branches && step.branches.length > 1) {
                    const branchResults = [];
                    const branchWidths = step.branches.map(b => this._calcStepWidth(b));
                    const totalWidth = branchWidths.reduce((a, b) => a + b, 0) + (step.branches.length - 1) * this.branchGap;

                    let branchX = parentCenterX - totalWidth / 2;
                    const branchStartY = y;

                    for (let bi = 0; bi < step.branches.length; bi++) {
                        const branch = step.branches[bi];
                        const branchWidth = branchWidths[bi];
                        const branchCenterX = branchX + branchWidth / 2;

                        const result = layoutSteps([branch], branchStartY, branchCenterX);
                        branchResults.push({
                            result,
                            centerX: branchCenterX,
                            firstNode: nodes.find(n => n.data === branch)
                        });

                        branchX += branchWidth + this.branchGap;
                    }

                    for (const prev of prevNodes) {
                        for (const br of branchResults) {
                            if (br.firstNode) {
                                arrows.push({ source: prev, target: br.firstNode });
                            }
                        }
                    }

                    let maxEndY = y;
                    let allLastNodes = [];
                    for (const br of branchResults) {
                        maxEndY = Math.max(maxEndY, br.result.endY);
                        if (br.result.lastNodes) {
                            allLastNodes.push(...br.result.lastNodes);
                        }
                    }

                    y = maxEndY;
                    prevNodes = allLastNodes;
                    continue;
                }

                const nodeWidth = this._getNodeWidth(step.params || 0);
                const nodeHeight = this._getNodeHeight(step.params || 0);
                const hasSubsteps = !!(step.substeps && step.substeps.length > 0);
                const isExpanded = hasSubsteps && step._collapsed === false;
                const x = parentCenterX - nodeWidth / 2;

                const node = {
                    data: step,
                    x: x,
                    y: y,
                    width: nodeWidth,
                    height: nodeHeight,
                    hasSubsteps: hasSubsteps,
                    collapsed: !isExpanded
                };

                nodes.push(node);

                for (const prev of prevNodes) {
                    arrows.push({ source: prev, target: node });
                }

                y += nodeHeight;

                if (isExpanded) {
                    const containerStartY = y + this.containerGap;
                    const childrenMaxWidth = Math.max(...step.substeps.map(s => this._calcStepWidth(s)));
                    const containerWidth = Math.max(nodeWidth, childrenMaxWidth) + this.containerPadding * 2;
                    const containerX = parentCenterX - containerWidth / 2;

                    const childResult = layoutSteps(
                        step.substeps,
                        containerStartY + this.containerPadding,
                        parentCenterX
                    );

                    const containerEndY = childResult.endY + this.containerPadding - this.nodeGap;

                    containers.push({
                        x: containerX,
                        y: containerStartY,
                        width: containerWidth,
                        height: containerEndY - containerStartY,
                        color: this.getColor(step.type)
                    });

                    const firstChild = nodes.find(n => n.data === step.substeps[0]);
                    if (firstChild) {
                        arrows.push({ source: node, target: firstChild });
                    }

                    y = containerEndY + this.containerGap;
                    prevNodes = [{
                        isContainerBottom: true,
                        containerBottomY: containerEndY,
                        centerX: parentCenterX
                    }];
                } else {
                    y += this.nodeGap;
                    prevNodes = [node];
                }
            }

            return { endY: y, lastNodes: prevNodes };
        };

        layoutSteps(data.steps || [], this.padding, centerX);

        return { nodes, arrows, containers };
    }

    _render() {
        if (!this.pipelineData) return;

        const { nodes, arrows, containers } = this._layoutPipeline(this.pipelineData);
        const arrowId = `arrow-${this.containerId}`;

        // Containers
        this.containersGroup.selectAll('.container-box').remove();
        containers.forEach(container => {
            this.containersGroup.append('rect')
                .attr('class', 'container-box')
                .attr('x', container.x)
                .attr('y', container.y)
                .attr('width', container.width)
                .attr('height', container.height)
                .attr('rx', 6)
                .attr('ry', 6)
                .attr('fill', 'rgba(24, 24, 27, 0.6)')
                .attr('stroke', container.color)
                .attr('stroke-width', 1)
                .attr('stroke-opacity', 0.3);
        });

        // Arrows
        this.arrowsGroup.selectAll('.arrow').remove();
        arrows.forEach(arrow => {
            let sx, sy;
            if (arrow.source.isContainerBottom) {
                sx = arrow.source.centerX;
                sy = arrow.source.containerBottomY;
            } else {
                sx = arrow.source.x + arrow.source.width / 2;
                sy = arrow.source.y + arrow.source.height;
            }

            const tx = arrow.target.x + arrow.target.width / 2;
            const ty = arrow.target.y;

            if (Math.abs(sx - tx) > 5) {
                const midY = (sy + ty) / 2;
                this.arrowsGroup.append('path')
                    .attr('class', 'arrow')
                    .attr('d', `M${sx},${sy} C${sx},${midY} ${tx},${midY} ${tx},${ty - 3}`)
                    .attr('fill', 'none')
                    .attr('stroke', '#3f3f46')
                    .attr('stroke-width', 1.5)
                    .attr('marker-end', `url(#${arrowId})`);
            } else {
                this.arrowsGroup.append('line')
                    .attr('class', 'arrow')
                    .attr('x1', sx)
                    .attr('y1', sy)
                    .attr('x2', tx)
                    .attr('y2', ty - 3)
                    .attr('stroke', '#3f3f46')
                    .attr('stroke-width', 1.5)
                    .attr('marker-end', `url(#${arrowId})`);
            }
        });

        // Nodes
        const nodeGroups = this.nodesGroup.selectAll('.node')
            .data(nodes, (d, i) => d.data.name + '-' + i)
            .join('g')
            .attr('class', 'node')
            .attr('transform', d => `translate(${d.x}, ${d.y})`)
            .style('cursor', d => d.hasSubsteps ? 'pointer' : 'default');

        nodeGroups.selectAll('rect')
            .data(d => [d])
            .join('rect')
            .attr('width', d => d.width)
            .attr('height', d => d.height)
            .attr('rx', 6)
            .attr('ry', 6)
            .attr('fill', d => this.getColor(d.data.type))
            .attr('stroke', d => d.hasSubsteps && d.collapsed ? 'rgba(255,255,255,0.25)' : 'none')
            .attr('stroke-width', 1)
            .attr('stroke-dasharray', d => d.hasSubsteps && d.collapsed ? '3,2' : 'none');

        // Name
        nodeGroups.selectAll('.node-name')
            .data(d => [d])
            .join('text')
            .attr('class', 'node-name')
            .attr('x', d => d.width / 2)
            .attr('y', d => d.data.shape ? 13 : (d.height / 2 - 3))
            .attr('text-anchor', 'middle')
            .attr('fill', 'white')
            .attr('font-size', '10px')
            .attr('font-weight', '600')
            .attr('font-family', '-apple-system, BlinkMacSystemFont, sans-serif')
            .text(d => {
                let name = d.data.name;
                if (d.data.count) name += ` (${d.data.count}x)`;
                return name;
            })
            .each(function(d) {
                const textEl = d3.select(this);
                let text = textEl.text();
                const maxWidth = d.width - 20;
                while (this.getComputedTextLength() > maxWidth && text.length > 0) {
                    text = text.slice(0, -1);
                    textEl.text(text + '...');
                }
            });

        // Shape
        nodeGroups.selectAll('.node-shape')
            .data(d => d.data.shape ? [d] : [])
            .join('text')
            .attr('class', 'node-shape')
            .attr('x', d => d.width / 2)
            .attr('y', d => d.height / 2)
            .attr('text-anchor', 'middle')
            .attr('fill', 'rgba(255,255,255,0.7)')
            .attr('font-size', '8px')
            .attr('font-family', 'SF Mono, monospace')
            .text(d => d.data.shape)
            .each(function(d) {
                const textEl = d3.select(this);
                let text = textEl.text();
                const maxWidth = d.width - 14;
                while (this.getComputedTextLength() > maxWidth && text.length > 0) {
                    text = text.slice(0, -1);
                    textEl.text(text + '...');
                }
            });

        // Params
        nodeGroups.selectAll('.node-params')
            .data(d => [d])
            .join('text')
            .attr('class', 'node-params')
            .attr('x', d => d.width / 2)
            .attr('y', d => d.data.shape ? d.height - 7 : (d.height / 2 + 10))
            .attr('text-anchor', 'middle')
            .attr('fill', 'rgba(255,255,255,0.5)')
            .attr('font-size', '9px')
            .attr('font-family', 'SF Mono, monospace')
            .text(d => API.formatParams(d.data.params || 0));

        // Expand indicator
        nodeGroups.selectAll('.expand-indicator')
            .data(d => d.hasSubsteps ? [d] : [])
            .join('text')
            .attr('class', 'expand-indicator')
            .attr('x', d => d.width - 10)
            .attr('y', 12)
            .attr('fill', 'rgba(255,255,255,0.5)')
            .attr('font-size', '8px')
            .text(d => d.collapsed ? '+' : '-');

        nodeGroups
            .on('click', (event, d) => this._handleClick(event, d))
            .on('mouseenter', (event, d) => this._handleMouseEnter(event, d))
            .on('mouseleave', (event, d) => this._handleMouseLeave(event, d));

        this._fitView(nodes, containers);
    }

    _fitView(nodes, containers) {
        if (nodes.length === 0) return;

        const rect = this.container.getBoundingClientRect();
        const padding = 32;

        let minX = Infinity, minY = Infinity, maxX = -Infinity, maxY = -Infinity;

        for (const node of nodes) {
            minX = Math.min(minX, node.x);
            minY = Math.min(minY, node.y);
            maxX = Math.max(maxX, node.x + node.width);
            maxY = Math.max(maxY, node.y + node.height);
        }

        for (const container of containers) {
            minX = Math.min(minX, container.x);
            minY = Math.min(minY, container.y);
            maxX = Math.max(maxX, container.x + container.width);
            maxY = Math.max(maxY, container.y + container.height);
        }

        const contentWidth = maxX - minX + padding * 2;
        const contentHeight = maxY - minY + padding * 2;

        const scaleX = rect.width / contentWidth;
        const scaleY = rect.height / contentHeight;
        const scale = Math.min(scaleX, scaleY, 1.5);

        const translateX = (rect.width - contentWidth * scale) / 2 - minX * scale + padding * scale;
        const translateY = (rect.height - contentHeight * scale) / 2 - minY * scale + padding * scale;

        this.svg.transition().duration(250).call(
            this.zoom.transform,
            d3.zoomIdentity.translate(translateX, translateY).scale(scale)
        );
    }

    _handleClick(event, d) {
        event.stopPropagation();
        if (d.hasSubsteps) {
            d.data._collapsed = d.data._collapsed === false ? true : false;
            this._render();
        }
    }

    _handleMouseEnter(event, d) {
        this._showTooltip(event, d);
    }

    _handleMouseLeave(event, d) {
        this._hideTooltip();
    }

    _showTooltip(event, d) {
        const tooltip = document.getElementById('tooltip');
        const params = d.data.params || 0;
        const refParams = this.globalMaxParams || this.totalParams;
        const percent = refParams > 0 ? ((params / refParams) * 100).toFixed(1) : '0';

        let extra = '';
        if (d.data.shape) {
            extra += `<div class="tooltip-row"><span class="tooltip-label">Shape</span><span class="tooltip-value">${d.data.shape}</span></div>`;
        }
        if (d.data.count) {
            extra += `<div class="tooltip-row"><span class="tooltip-label">Layers</span><span class="tooltip-value">${d.data.count}</span></div>`;
        }
        if (d.data.substeps && d.data.substeps.length > 0) {
            extra += `<div class="tooltip-row"><span class="tooltip-label">Components</span><span class="tooltip-value">${d.data.substeps.length}</span></div>`;
        }

        tooltip.innerHTML = `
            <div class="tooltip-title">${d.data.name}</div>
            <div class="tooltip-content">
                <div class="tooltip-row"><span class="tooltip-label">Parameters</span><span class="tooltip-value">${API.formatParams(params)}</span></div>
                <div class="tooltip-row"><span class="tooltip-label">Proportion</span><span class="tooltip-value">${percent}%</span></div>
                <div class="tooltip-row"><span class="tooltip-label">Type</span><span class="tooltip-value">${d.data.type}</span></div>
                ${extra}
            </div>
        `;

        tooltip.classList.remove('hidden');

        const pad = 10;
        let x = event.clientX + pad;
        let y = event.clientY + pad;

        const tooltipRect = tooltip.getBoundingClientRect();
        if (x + tooltipRect.width > window.innerWidth) x = event.clientX - tooltipRect.width - pad;
        if (y + tooltipRect.height > window.innerHeight) y = event.clientY - tooltipRect.height - pad;

        tooltip.style.left = `${x}px`;
        tooltip.style.top = `${y}px`;
    }

    _hideTooltip() {
        document.getElementById('tooltip').classList.add('hidden');
    }

    getLegendItems() {
        if (!this.pipelineData || !this.pipelineData.steps) return [];

        const types = new Set();
        const collect = (steps) => {
            if (!steps) return;
            for (const step of steps) {
                if (step.type) types.add(step.type);
                if (step.substeps) collect(step.substeps);
                if (step.branches) {
                    for (const branch of step.branches) {
                        if (branch.type) types.add(branch.type);
                        if (branch.substeps) collect(branch.substeps);
                    }
                }
            }
        };
        collect(this.pipelineData.steps);

        return Array.from(types).map(type => ({ type, color: this.getColor(type) }));
    }

    /**
     * Export SVG with watermark
     */
    exportSVG(modelName) {
        const svgElement = this.svg.node();
        const clone = svgElement.cloneNode(true);

        // Get bounds
        const rect = this.container.getBoundingClientRect();
        clone.setAttribute('width', rect.width);
        clone.setAttribute('height', rect.height);
        clone.setAttribute('viewBox', `0 0 ${rect.width} ${rect.height}`);

        // Add background
        const bg = document.createElementNS('http://www.w3.org/2000/svg', 'rect');
        bg.setAttribute('width', '100%');
        bg.setAttribute('height', '100%');
        bg.setAttribute('fill', '#0a0a0b');
        clone.insertBefore(bg, clone.firstChild);

        // Add watermark
        const watermark = document.createElementNS('http://www.w3.org/2000/svg', 'text');
        watermark.setAttribute('x', rect.width - 10);
        watermark.setAttribute('y', rect.height - 10);
        watermark.setAttribute('text-anchor', 'end');
        watermark.setAttribute('fill', '#3f3f46');
        watermark.setAttribute('font-size', '11');
        watermark.setAttribute('font-family', '-apple-system, sans-serif');
        watermark.textContent = 'omarkamali.com/llm-scope';
        clone.appendChild(watermark);

        // Add model name
        if (modelName) {
            const title = document.createElementNS('http://www.w3.org/2000/svg', 'text');
            title.setAttribute('x', '10');
            title.setAttribute('y', '20');
            title.setAttribute('fill', '#fafafa');
            title.setAttribute('font-size', '12');
            title.setAttribute('font-weight', '600');
            title.setAttribute('font-family', '-apple-system, sans-serif');
            title.textContent = modelName;
            clone.appendChild(title);
        }

        const serializer = new XMLSerializer();
        const svgString = serializer.serializeToString(clone);
        const blob = new Blob([svgString], { type: 'image/svg+xml' });

        return blob;
    }

    resetZoom() {
        if (this.pipelineData) {
            const expandAll = (steps) => {
                if (!steps) return;
                for (const step of steps) {
                    step._collapsed = false;
                    if (step.substeps) expandAll(step.substeps);
                    if (step.branches) {
                        for (const branch of step.branches) {
                            branch._collapsed = false;
                            if (branch.substeps) expandAll(branch.substeps);
                        }
                    }
                }
            };
            if (this.pipelineData.steps) {
                expandAll(this.pipelineData.steps);
            }
            this._render();
        }
    }

    destroy() {
        if (this.resizeObserver) this.resizeObserver.disconnect();
        this.svg.remove();
    }
}

const TreemapViz = ModelTreeViz;