document.addEventListener("DOMContentLoaded", function() {
    const select = document.getElementById('language-family-select');
    const hoverBox = document.getElementById('hover-box');
    const searchBar = document.getElementById('search-bar');
    const depthSelector = document.getElementById('depth-number');
    const tokenizerFilterDiv = document.getElementById('tokenizer-filter');
    const clearTokenizersFilter = document.getElementById('clear-tokenizers');
    let nodeToCenter = null;
    let currentTreeData = null;
    let expandedNodes;
    let currentTransform = d3.zoomIdentity; // Save the current zoom and transform
        // Get unique tokenizer names
    let tokenizerNames;
    let color;

    function setExpanded(id, state = true){
        if (state) {
            expandedNodes.add(id)
        }
        else if(expandedNodes.has(id)) {
            expandedNodes.delete(id)
        }
    }

    function addNodeIds(tree){
        function addIdToSubtree(tree, nodeCount = 0){
            tree.id = nodeCount++;
            let tree_subtreeSize = tree.children.length === 0 ? 1 : 0;
            let node = {
                ...tree,
                children: (tree.children || []).map(child => {
                    const [subtree, newCount, subtreeSize] = addIdToSubtree(child, nodeCount);
                    nodeCount = newCount;
                    tree_subtreeSize += subtreeSize;
                    return subtree;
                })
            };
            node.subtreeSize = tree_subtreeSize
            return [node, nodeCount, tree_subtreeSize];
        }
        const [parsedTree, _finalCount] = addIdToSubtree(tree);
        return parsedTree;
    }

    function loadLanguageFamily() {
        const family = select.value;
        fetch(`data/${family}.json`)
            .then(response => response.json())
            .then(data => {
                currentTreeData = addNodeIds(data);
                expandedNodes = new Set([0]);
                updateTokenizerFilter(data);
                drawVisibleNodes(true);
            });
    }
    loadLanguageFamily();
    select.addEventListener('change', loadLanguageFamily);
    searchBar.addEventListener('input', () => searchNode(searchBar.value));

    depthSelector.addEventListener('change', (_) => drawVisibleNodes());
    clearTokenizersFilter.addEventListener('click', (_) => {
        tokenizerFilterDiv.querySelectorAll('input:checked').forEach((a) => a.checked = false);
        drawVisibleNodes();
    })

    function updateTokenizerFilter(data) {
        tokenizerNames = Array.from(getTokenizerNames(data));
        // Create color scale based on the unique tokenizer names
        color = d3.scaleOrdinal(tokenizerNames, d3.schemeCategory10);
        tokenizerFilterDiv.innerHTML = '';
        tokenizerNames.forEach(name => {
            const checkbox = document.createElement('input');
            checkbox.type = 'checkbox';
            checkbox.value = name;
            checkbox.checked = true;
            checkbox.addEventListener('change', () => drawVisibleNodes());
            const label = document.createElement('label');
            // Create the colored icon
            const icon = document.createElement('span');
            icon.style.display = 'inline-block';
            icon.style.width = '10px';
            icon.style.height = '10px';
            icon.style.backgroundColor = color(name);
            icon.style.marginRight = '5px';

            label.appendChild(icon); // Add the icon before the checkbox
            label.appendChild(checkbox);
            label.appendChild(document.createTextNode(name));

            tokenizerFilterDiv.appendChild(label);
            tokenizerFilterDiv.appendChild(document.createElement('br'));
        });
    }

    function getTokenizerNames(node, namesSet = new Set()) {
        if (!node)
            return namesSet;
        if (node.tokenizer && node.tokenizer.name) {
            namesSet.add(node.tokenizer.name);
        } else
            namesSet.add('x')
        if (node.children) {
            node.children.forEach(child => getTokenizerNames(child, namesSet));
        }
        return namesSet;
    }

    function createTree(data, recenter = false) {
        // Clear any existing tree
        d3.select("#tree-container svg").remove();

        // Set the SVG dimensions to fill the entire screen
        const width = window.innerWidth;
        const height = window.innerHeight;

        // Create the SVG element
        const svg = d3.select("#tree-container")
            .append("svg")
            .attr("width", width)
            .attr("height", height);

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

        // Define zoom behavior without restricting translation extents
        const zoom = d3.zoom()
            .scaleExtent([0.1, 5])
            .on("zoom", function(event) {
                currentTransform = event.transform; // Save the current transform on zoom
                g.attr("transform", event.transform);
            });

        // Create a root hierarchy node
        const root = d3.hierarchy(data);

        // Dynamically compute tree layout for horizontal orientation
        const treeLayout = d3.tree().nodeSize([200, 100]);
        treeLayout(root);

        // Create links between nodes (now horizontal)
        g.selectAll('.link')
            .data(root.links())
            .enter()
            .append('path')
            .attr('class', 'link')
            .attr('d', d3.linkVertical()
                .x(d => d.x)
                .y(d => d.y))
            .attr('stroke', '#ccc')
            .attr('fill', 'none');

        // Create nodes
        const node = g.selectAll('.node')
            .data(root.descendants())
            .enter()
            .append('g')
            .attr('class', 'node')
            .attr('transform', d => `translate(${d.x},${d.y})`);

        const sizeScale = d3.scaleSqrt()
            .domain([1, root.data.subtreeSize])
            .range([5, 20]); // Adjust the range as needed for minimum and maximum circle

        // Add circles to nodes
        node.append('circle')
            .filter(d => !d.data.iso_3_code)
            .attr('r', d => sizeScale(d.data.subtreeSize))
            .attr('fill', d => {
                const tokenizerName = d.data.tokenizer?.name || "unknown";
                return color(tokenizerName);
            });
        node.filter(d => d.data.iso_3_code && d.data.source !== 'own' && d.data.source !== 'macrolanguage') // Select leaf nodes
            .append('rect')
            .attr('width', 10)
            .attr('height', 10)
            .attr('x', -5)
            .attr('y', -5)
            .attr('fill', d => {
                const tokenizerName = d.data.tokenizer?.name || "unknown";
                return color(tokenizerName);
            });
        node.filter(d => d.data.source === 'own' || d.data.source === 'macrolanguage') // Select leaf nodes with "own" assignment
            .append('path')
            .attr('d', d3.symbol().type(d3.symbolTriangle).size(100)) // Adjust size as needed
            .attr('fill', d => {
                const tokenizerName = d.data.tokenizer?.name || "unknown";
                return color(tokenizerName);
            });

        // Add text labels to nodes
        node.append('text')
            .attr('dy', 4)
            .attr('x', d => sizeScale(d.data.subtreeSize) + 4)
            .attr('text-anchor', 'start')
            .text(d => `${d.data.name} - ${d.data.tokenizer?.name || 'x'}${d.data.iso_3_code ? '' : ' (' + d.data.subtreeSize + ')'}`);

        if (!currentTransform || recenter)
            currentTransform = d3.zoomIdentity.translate(width / 2, height / 2);

        if (nodeToCenter){
            debugger;
            let element = g.selectAll('.node').filter(d => d.data.id === nodeToCenter).data()[0];
            if (element){
                const x = element.x;
                const y = element.y;
                const scale = 1.0; // Adjust scale as needed

                currentTransform = d3.zoomIdentity.translate(window.innerWidth / 2 - x, window.innerHeight / 2 - y).scale(scale);
            }
            nodeToCenter = null;
        }
        svg.call(zoom).call(zoom).call(zoom.transform, d3.zoomIdentity.translate(currentTransform.x, currentTransform.y).scale(currentTransform.k));


        // Add hover event to nodes for hover-box
        node.on("mouseover", function(event, d) {
            hoverBox.style.display = "block";
            hoverBox.style.left = (event.pageX) + "px";
            hoverBox.style.top = (event.pageY) + "px";
            hoverBox.innerHTML = `
                <strong>Name:</strong> ${d.data.name}<br>
                <strong>ISO 1 Code:</strong> ${d.data.iso_1_code}<br>
                <strong>ISO 3 Code:</strong> ${d.data.iso_3_code}<br>
                <strong>Tokenizer:</strong> ${d.data.tokenizer?.tokenizer} (${d.data.tokenizer?.name})<br>
                <strong>Tokenizer assignment:</strong> ${d.data.source} <br>
                <strong>Subtree size:</strong> ${d.data.subtreeSize}
            `;
        }).on("mousemove", function(event) {
            hoverBox.style.left = (event.pageX) + "px";
            hoverBox.style.top = (event.pageY) + "px";
        }).on("mouseout", function() {
            hoverBox.style.display = "none";
        });

        node.on('click', function(event, d) {
            setExpanded(d.data.id, !expandedNodes.has(d.data.id))
            hoverBox.style.display = "none";
            drawVisibleNodes();
        })
    }

    function drawVisibleNodes(recenter = false) {
        if (!currentTreeData) return;

        // Get selected tokenizers
        const selectedTokenizers = Array.from(tokenizerFilterDiv.querySelectorAll('input:checked'))
            .map(input => input.value);

        // Create a new root node containing only selected tokenizers and their ancestors
        function filterHierarchy(node, parentExpanded = true, depth = 0) {
            if (!parentExpanded && !(depthSelector.value == 0 || depth <= depthSelector.value)) {
                return null;
            }

            const filteredChildren = (node.children || [])
                    .map(child => filterHierarchy(child, expandedNodes.has(node.id), depth + 1))
                    .filter(child => child !== null);

            if (!parentExpanded && filteredChildren.length === 0 && (!selectedTokenizers.includes(node.tokenizer ? node.tokenizer.name : 'x')) && node.id != 0)
                return null;

            return {
                ...node,
                children: filteredChildren
            };
        }

        const filteredData = filterHierarchy(currentTreeData);
        createTree(filteredData, recenter);
    }

    function searchNode(name) {
        name = name.toLowerCase();
        if (!currentTreeData || name.length < 2) return;

        // Find node by name
        function expandNode(node) {
            if (node.name.toLowerCase() === name || node.iso_3_code === name || node.iso_1_code === name) {
                return node;
            }

            for (const child of (node.children || [])) {
                const found = expandNode(child);
                if (found) {
                    setExpanded(node.id)
                    return found;
                }
            }

            return null;
        }

        const expandedNode = expandNode(currentTreeData);
        if (expandedNode) {
            nodeToCenter = expandedNode.id;
            drawVisibleNodes();
        }
    }
});