Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h

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

#pragma once

#include "BraggIntegrationEngine.h"
#include "../../common/PixelMask.h"

class CompressedImage;

// Plain-C++ reference/fallback engine: a faithful serial re-expression of BraggIntegrate2D (box
// sum) and ProfileIntegrate2D (Kabsch profile fit) reading the preprocessed int32 image. Also the
// numeric oracle the CUDA engine is checked against.
class BraggIntegrationEngineCPU : public BraggIntegrationEngine {
    // Core integrator, templated on a pixel sampler so it reads either the preprocessed int32 buffer
    // or a raw CompressedImage of any pixel type - both presented per-pixel in the INT32_MIN(masked)/
    // INT32_MAX(saturated) convention - without ever materialising a second full-image copy.
    // Full-frame scratch of RunImpl: the reflection mask and the signal-region owner map. A call writes
    // only around its reflections, so these keep only the 16x16-pixel tiles written since the last
    // Clear(): a frame-sized array was mostly never read, yet over a sweep every page of it got
    // touched, in every worker. Reading a tile nothing wrote gives `empty`.
    template <class T>
    class TiledFrame {
        static constexpr int TILE = 16;
        int tiles_x;
        T empty;
        std::vector<int32_t> tile_start;   // per tile: where it starts in `pixels`, -1 = not written
        std::vector<int32_t> written;      // the tiles written, for Clear()
        std::vector<T> pixels;
    public:
        TiledFrame(int width, int height, T empty)
            : tiles_x((width + TILE - 1) / TILE), empty(empty),
              tile_start(static_cast<size_t>(tiles_x) * ((height + TILE - 1) / TILE), -1) {}
        T Get(int x, int y) const {
            const int32_t start = tile_start[(y / TILE) * tiles_x + x / TILE];
            return start < 0 ? empty : pixels[start + (y % TILE) * TILE + x % TILE];
        }
        T &At(int x, int y) {
            const int t = (y / TILE) * tiles_x + x / TILE;
            if (tile_start[t] < 0) {
                tile_start[t] = static_cast<int32_t>(pixels.size());
                pixels.resize(pixels.size() + TILE * TILE, empty);
                written.push_back(t);
            }
            return pixels[tile_start[t] + (y % TILE) * TILE + x % TILE];
        }
        void Clear() {
            for (int t : written)
                tile_start[t] = -1;
            written.clear();
            pixels.clear();
        }
    };
    TiledFrame<uint8_t> refl_mask;
    TiledFrame<uint32_t> owner;
    // The run's pixel mask, packed 32 pixels to a word (PixelMask::GetPackedMask). An unreadable pixel
    // it does not explain was unreadable on this frame only - an overload (see Reflection::overloaded).
    std::vector<uint32_t> static_mask;

    // Marks every predicted reflection's r2 region in refl_mask (MASK_FLUX / MASK_TAIL).
    void MarkNeighbours(const std::vector<Reflection> &predicted, size_t npredicted);

    template <class Sampler>
    std::vector<Reflection> RunImpl(const Sampler &img, const std::vector<Reflection> &predicted,
                                    size_t npredicted, int64_t image_number);

public:
    BraggIntegrationEngineCPU(const DiffractionExperiment &experiment, const PixelMask &mask);

    using BraggIntegrationEngine::Run;   // keep the preprocessed-buffer overload visible

    std::vector<Reflection> Run(const ImagePreprocessorBuffer &image,
                                const std::vector<Reflection> &predicted, size_t npredicted,
                                int64_t image_number) override;

    // Only the neighbour count of these predicted reflections at this engine's stencil
    // (BraggIntegrationCounts::predicted and bkg_starved_by_neighbour), reading no image: what a pass
    // integrating at one radius measures of another. A ring pixel counts as readable unless the run's
    // mask excludes it, so an overload on one frame is read as readable here.
    void CountNeighbourCrowding(const std::vector<Reflection> &predicted, size_t npredicted);

    // FPGA workflow: integrate straight off the assembled detector image, reading only the pixels
    // inside each reflection disk (no whole-image conversion - the FPGA host cannot afford one at its
    // frame rate). Masked pixels carry the type minimum and saturated the type maximum.
    std::vector<Reflection> Run(const CompressedImage &image,
                                const std::vector<Reflection> &predicted, size_t npredicted,
                                int64_t image_number);
};