// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "ConnectedComponents.h"

std::vector<int32_t> LabelConnectedComponents(const std::vector<uint8_t> &map,
                                              int64_t nx, int64_t ny, int64_t min_size) {
    std::vector<int32_t> label(map.size(), 0);
    std::vector<int64_t> stack;    // cells of the component being grown
    std::vector<int64_t> members;  // same, kept so the component can be erased if it is too small
    int32_t next_label = 0;

    for (int64_t seed = 0; seed < nx * ny; seed++) {
        if (!map[seed] || label[seed] != 0)
            continue;

        next_label++;
        stack = {seed};
        members.clear();
        label[seed] = next_label;

        while (!stack.empty()) {
            int64_t i = stack.back();
            stack.pop_back();
            members.push_back(i);

            const int64_t x = i % nx;
            const int64_t y = i / nx;
            const bool left = x > 0, right = x < nx - 1, up = y > 0, down = y < ny - 1;
            const int64_t neighbour[8] = {left ? i - 1 : -1,
                                          right ? i + 1 : -1,
                                          up ? i - nx : -1,
                                          down ? i + nx : -1,
                                          left && up ? i - nx - 1 : -1,
                                          right && up ? i - nx + 1 : -1,
                                          left && down ? i + nx - 1 : -1,
                                          right && down ? i + nx + 1 : -1};
            for (int64_t j : neighbour) {
                if (j >= 0 && map[j] && label[j] == 0) {
                    label[j] = next_label;
                    stack.push_back(j);
                }
            }
        }

        if (static_cast<int64_t>(members.size()) < min_size) {
            for (int64_t i : members)
                label[i] = 0;
        }
    }

    return label;
}