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

#include <bitset>

#include "ImageSpotFinder.h"
#include "StrongPixelSet.h"

ImageSpotFinder::ImageSpotFinder(int32_t width, int32_t height)
    : width(width),
      height(height),
      output_buffer(width * height / 32 + 1) {
}

size_t ImageSpotFinder::OutputSize() const {
    return (width * height) / 32 + ((width * height % 32 != 0) ? 1 : 0);
}

size_t ImageSpotFinder::OutputByteSize() const {
    return OutputSize() * sizeof(uint32_t);
}

std::vector<DiffractionSpot> ImageSpotFinder::Run(const ImagePreprocessorBuffer &image,
                                                 const SpotFindingSettings &settings,
                                                 const std::vector<bool> &res_mask) {
    Detect(image, settings);
    return ExtractSpots(image, settings, res_mask);
}

std::vector<DiffractionSpot> ImageSpotFinder::ExtractSpots(const ImagePreprocessorBuffer &image,
                                                           const SpotFindingSettings &settings,
                                                           const std::vector<bool> &res_mask) {
    // Collect the strong pixels first and read their values afterwards, instead of reading the image
    // pixel by pixel: on the GPU that read is a device gather, which is what lets the preprocessed
    // image stay on the device instead of being copied back in full for every frame.
    strong_pixel.clear();
    for (int i = 0; i < OutputSize(); i++) {
        if (output_buffer[i]) {
            std::bitset<32> bset = output_buffer[i];
            for (int bit = 0; bit < 32; bit++) {
                if (bset.test(bit)) {
                    uint32_t npixel = i * 32 + bit;
                    if (npixel / width < height && res_mask[npixel] == 0)
                        strong_pixel.push_back(npixel);
                }
            }
        }
    }

    std::vector<DiffractionSpot> vec;
    // The connected-component search rejects a frame with this many strong pixels, so their values are
    // of no use.
    if (strong_pixel.size() >= UINT16_MAX)
        return vec;

    image.Gather(strong_pixel, strong_pixel_value);

    StrongPixelSet pixel_set;
    for (size_t i = 0; i < strong_pixel.size(); i++)
        pixel_set.AddStrongPixel(strong_pixel[i] % width, strong_pixel[i] / width, strong_pixel_value[i]);
    pixel_set.FindSpotsImage(settings, vec);
    return vec;
}