export function connectedComponentLabeling(binaryImage: ArrayLike<number>, width: number, height: number) {
    const labels: number[] = Array(binaryImage.length).fill(0);
    const linked: number[] = [];
    let nextLabel = 1;

    function find(x: number): number {
        if (x !== linked[x]) {
            linked[x] = find(linked[x]);
        }
        return linked[x];
    }

    function union(x: number, y: number) {
        linked[find(x)] = find(y);
    }

    function getNeighbors(row: number, col: number) {
        const neighbors: number[] = [];
        // Above
        if (row > 0 && labels[(row - 1) * width + col] > 0) {
            neighbors.push(labels[(row - 1) * width + col]);
        }
        // Left
        if (col > 0 && labels[row * width + col - 1] > 0) {
            neighbors.push(labels[row * width + col - 1]);
        }
        // Above Left
        if (row > 0 && col > 0 && labels[(row - 1) * width + col - 1] > 0) {
            neighbors.push(labels[(row - 1) * width + col - 1]);
        }
        // Above Right
        if (row > 0 && col < width - 1 && labels[(row - 1) * width + col + 1] > 0) {
            neighbors.push(labels[(row - 1) * width + col + 1]);
        }
        return neighbors;
    }

    // First pass
    for (let row = 0; row < height; row++) {
        for (let col = 0; col < width; col++) {
            const idx = row * width + col;
            if (binaryImage[idx] !== 0) {
                const neighbors = getNeighbors(row, col);

                if (neighbors.length === 0) {
                    linked[nextLabel] = nextLabel;
                    labels[idx] = nextLabel;
                    nextLabel++;
                } else {
                    neighbors.sort();
                    const smallestLabel = neighbors[0];
                    labels[idx] = smallestLabel;

                    for (const neighbor of neighbors) {
                        if (neighbor !== smallestLabel) {
                            union(smallestLabel, neighbor);
                        }
                    }
                }
            }
        }
    }

    // Second pass
    for (let idx = 0; idx < binaryImage.length; idx++) {
        if (binaryImage[idx] !== 0) {
            labels[idx] = find(labels[idx]);
        }
    }

    return labels;
}


export function computeBounds(labels: number[], width: number, height: number) {
    const bounds: Record<number, { minRow: number, minCol: number, maxRow: number, maxCol: number; area: number; }> = {};

    for (let row = 0; row < height; row++) {
        for (let col = 0; col < width; col++) {
            const idx = row * width + col;
            const label = labels[idx];

            if (label > 0) {
                if (!bounds[label]) {
                    bounds[label] = { minRow: row, minCol: col, maxRow: row, maxCol: col, area: 1 };
                } else {
                    if (row < bounds[label].minRow) bounds[label].minRow = row;
                    if (col < bounds[label].minCol) bounds[label].minCol = col;
                    if (row > bounds[label].maxRow) bounds[label].maxRow = row;
                    if (col > bounds[label].maxCol) bounds[label].maxCol = col;
                    ++bounds[label].area;
                }
            }
        }
    }

    return bounds;
}