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

#pragma once

#include <cstdint>
#include <vector>

#include "../../common/DiffractionSpot.h"

#include "../image_preprocessing/ImagePreprocessorBuffer.h"

class ImageSpotFinder {
     // Flat index of every strong pixel of the current image that passed the resolution mask, and its
     // value. Kept as members only to reuse the allocation from image to image.
     std::vector<uint32_t> strong_pixel;
     std::vector<int32_t> strong_pixel_value;
protected:
     const int32_t width, height;
     std::vector<uint32_t> output_buffer;
     // Pixels excluded from spot finding, packed the same way as output_buffer (bit set = excluded).
     // The bits past the last image pixel are set at construction, so the padding of the last word
     // needs no separate guard. Default: nothing excluded.
     std::vector<uint32_t> res_mask_bits;
     // The connected components of the last extraction. A member so ExtractComponents can hand out a
     // reference and reuse the allocation from image to image.
     std::vector<DiffractionSpot> components;

     // host_bit_buffer = false leaves output_buffer empty: the GPU finders extract on the device and
     // never read the bit buffer on the host, so allocating and pinning 2.26 MB per engine (at 18 MP)
     // would be pure waste for them.
     ImageSpotFinder(int32_t width, int32_t height, bool host_bit_buffer = true);
     size_t OutputSize() const;
     size_t OutputByteSize() const;
     // Host extraction: scan the bit buffer, gather the values, run the connected-component search.
     void ExtractComponentsHost(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings);
public:
     constexpr static int32_t MIN_VALID_PIXELS = 100;
     constexpr static int NBX = 15;

     virtual ~ImageSpotFinder() = default;

     // Detect flags the image's strong pixels into the internal bit buffer - the expensive step (local
     // box or per-ring background over every pixel). ExtractComponents then builds the connected
     // components from those pixels.
     virtual void Detect(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings) = 0;

     // Peak-excluded per-ring background of the last Detect(), in the bins of the azimuthal-integration
     // mapping and in raw photon counts. Only the adaptive finders build one (it is what sets their
     // threshold); empty for everyone else, and for a frame with nothing valid to reduce.
     [[nodiscard]] virtual const std::vector<float> &GetRingBackground() const;

     // Pixels to ignore, one bool per pixel (true = ignore). Set when the resolution limits change,
     // not per image: the GPU finders keep a bit-packed device copy of it, and re-uploading that for
     // every image would cost more than the extraction it feeds.
     virtual void SetResolutionMask(const std::vector<bool> &mask);

     // Every connected component of the last Detect() with at most max-pix pixels. min-pix is NOT
     // applied here on purpose - it is the only spot setting that changes between the passes of the
     // per-image min-pix search, so ONE extraction serves all three of them.
     virtual const std::vector<DiffractionSpot> &ExtractComponents(const ImagePreprocessorBuffer &image,
                                                                   const SpotFindingSettings &settings);
     // The components that also pass min-pix.
     static std::vector<DiffractionSpot> Filter(const std::vector<DiffractionSpot> &in,
                                                const SpotFindingSettings &settings);

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