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 "BraggPassA.h"
#include "BraggPassB.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, and the engine the Metal engine extends.
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;

    // Marks every predicted reflection's r2 region in refl_mask (MASK_FLUX / MASK_TAIL), visiting them in
    // `order` (any order gives the same mask).
    void MarkNeighbours(const std::vector<Reflection> &predicted, const std::vector<uint32_t> &order);

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

    // Pass B's single-precision part for reflection i, as the Metal engine's kernel does it: the profile
    // (the shell's learned one, or the Gaussian, built in `cells`), how much of it is readable, and the
    // fit. fit_pixels is scratch: the pixels the fit reads, gathered once.
    struct FitPixel { float Pp, signal; };
    template <class Sampler, class OwnerMap>
    BraggFitOut FitProfile(const Sampler &img, const OwnerMap &owner_map, const BraggFitSetup &s, size_t i,
                           const std::vector<float> &shell_profiles, std::vector<float> &cells,
                           std::vector<FitPixel> &fit_pixels) const;
    // Whether enough of a reflection's profile is readable, and its own, for the fit to be kept.
    bool FitKept(const BraggFitOut &o) const;

protected:
    // 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;

    // Identity sampler over the preprocessed int32 buffer (already in the engine's convention).
    struct BufferSampler {
        const int32_t *p;
        int32_t operator[](size_t i) const { return p[i]; }
        const void *address(size_t i) const { return p + i; }
    };
    // The owner map as one frame-sized array, as the Metal engine builds it (TiledFrame's Get otherwise).
    struct FlatOwnerMap {
        const uint32_t *p;
        size_t width;
        uint32_t Get(int x, int y) const { return p[static_cast<size_t>(y) * width + x]; }
    };

    // Pass A's result for one reflection: the box sum, its background, centroid and strong flag.
    struct Rough {
        double I = 0.0, sigma = NAN, bkg = 0.0, obs_x = 0.0, obs_y = 0.0;
        double bkg_var = 0.0;   // variance of the background ESTIMATE itself, bkg / n_bkg
        double var_bkg = 0.0;   // total non-signal variance carried to the merge
        int64_t I_sum = 0;      // kept so I can be rebuilt after the radial background correction
        int n_inner = 0;
        int n_disk = 0, n_own = 0;   // signal-disk pixels, and how many of them are this reflection's
        int r_bin = 0;          // rounded distance from the beam centre, indexes the radial curve
        int k_bin = 0;          // which radial-background kernel this reflection's stencil needs
        int cx = 0, cy = 0, shell = -1;
        bool ok = false, strong = false, has_obs = false;
        bool full = false;      // every pixel of the signal disk was readable
        bool overloaded = false;   // a pixel of the signal disk was saturated
    };
    // Everything pass A hands the profile passes.
    struct PassA {
        std::vector<Rough> rough;
        // Radial background curve, accumulated from the ring pixels the background clip keeps. Integer
        // sums, so the order the pixels arrive in does not matter.
        int n_rad = 0;
        std::vector<int64_t> rad_sum;
        std::vector<int> rad_cnt;
        double inv_d2_min = 0.0, inv_d2_max = 0.0;   // resolution range of the complete reflections
    };
    PassA StartPassA(size_t npredicted) const;

    // Whether pass A sigma-clips this reflection's background ring, and the clip threshold: the ring
    // pixels at or below it are kept. A pixel value is an integer, so the threshold is one too.
    bool ClipsRing(const BraggBoxSums &s) const;
    int32_t ClipThreshold(const BraggBoxSums &s) const;
    // Pass A's result for reflection i from what its box walk counted, and from its clip sums where
    // ClipsRing. bkg_vals holds the ring pixels for the trimmed-mean background, and is sorted here.
    void FinishPassA(PassA &a, size_t i, const Reflection &r, const BraggStencil &st, const BraggBoxSums &s,
                     const BraggClipSums &clip, std::vector<int32_t> &bkg_vals);
    // Everything after pass A: the radial background correction, then either the box-sum results or
    // the profile learning and pass B. Defined for BufferSampler with either owner map, and for the
    // CompressedImage samplers with the tiled one.
    // Pass B's single-precision part for every reflection with setups[i].fit, on a device - the
    // learned profiles side by side in shell_profiles, empty for the Gaussian. False where there is no
    // device, and the CPU fits them instead.
    virtual bool FitOnDevice(const std::vector<BraggFitSetup> &setups, const std::vector<float> &shell_profiles,
                             std::vector<BraggFitOut> &fits) { return false; }

    template <class Sampler, class OwnerMap>
    std::vector<Reflection> ProfilePasses(const Sampler &img, const OwnerMap &owner_map,
                                          const std::vector<Reflection> &predicted, size_t npredicted,
                                          PassA &a, 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);
};