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

#include <algorithm>
#include <bit>

#include "../../common/JFJochException.h"
#include "ImageSpotFinder.h"
#include "StrongPixelSet.h"

ImageSpotFinder::ImageSpotFinder(int32_t width, int32_t height, bool host_bit_buffer)
    : width(width),
      height(height),
      output_buffer(host_bit_buffer ? width * height / 32 + 1 : 0),
      res_mask_bits(OutputSize(), 0) {
    // Exclude the padding bits of the last word up front, so neither the host scan nor the GPU
    // compaction needs a separate "is this bit still inside the image?" test.
    const size_t npixel = static_cast<size_t>(width) * height;
    if (npixel % 32 != 0)
        res_mask_bits.back() = ~((1u << (npixel % 32)) - 1u);
}

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);
}

void ImageSpotFinder::SetResolutionMask(const std::vector<bool> &mask) {
    const size_t npixel = static_cast<size_t>(width) * height;
    if (mask.size() != npixel)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "ImageSpotFinder::SetResolutionMask: mask size mismatch");
    std::fill(res_mask_bits.begin(), res_mask_bits.end(), 0);
    for (size_t i = 0; i < npixel; i++)
        if (mask[i])
            res_mask_bits[i / 32] |= 1u << (i % 32);
    if (npixel % 32 != 0)
        res_mask_bits.back() |= ~((1u << (npixel % 32)) - 1u);
}

void ImageSpotFinder::ExtractComponentsHost(const ImagePreprocessorBuffer &image,
                                            const SpotFindingSettings &settings) {
    // 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 (size_t i = 0; i < OutputSize(); i++) {
        // The resolution mask is packed like the bit buffer, so a whole word of it is excluded at
        // once instead of testing 32 bits one at a time.
        uint32_t word = output_buffer[i] & ~res_mask_bits[i];
        while (word != 0) {
            strong_pixel.push_back(static_cast<uint32_t>(i * 32 + std::countr_zero(word)));
            word &= word - 1;
        }
    }

    components.clear();
    // 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;

    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.FindComponentsImage(settings, components);
}

const std::vector<DiffractionSpot> &ImageSpotFinder::ExtractComponents(const ImagePreprocessorBuffer &image,
                                                                       const SpotFindingSettings &settings) {
    ExtractComponentsHost(image, settings);
    return components;
}

std::vector<DiffractionSpot> ImageSpotFinder::Filter(const std::vector<DiffractionSpot> &in,
                                                     const SpotFindingSettings &settings) {
    std::vector<DiffractionSpot> out;
    const int64_t min_pix = settings.min_pix_per_spot.value_or(2);
    for (const auto &spot: in)
        if (spot.PixelCount() >= min_pix)
            out.push_back(spot);
    return out;
}

std::vector<DiffractionSpot> ImageSpotFinder::ExtractSpots(const ImagePreprocessorBuffer &image,
                                                           const SpotFindingSettings &settings) {
    return Filter(ExtractComponents(image, settings), settings);
}

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