Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngine.h

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

#pragma once

// =============================================================================
// BraggIntegrationEngine — box-sum + profile-fitting 2D integrator, GPU-ready
// =============================================================================
//
// A reimplementation of BraggIntegrate2D (box sum) and ProfileIntegrate2D (Kabsch profile
// fit) under one roof, following the AzIntEngine / ROIIntegration pattern: a base class that
// extracts the fixed per-experiment configuration, a plain-C++ CPU engine (the fallback and the
// numeric oracle), and a CUDA engine (BraggIntegrationEngineGPU) that reaches the same result up
// to floating-point precision.
//
// Unlike BraggIntegrate2D/ProfileIntegrate2D, which read the raw CompressedImage per pixel type
// and reject the special/saturation +/-1 band, this engine reads the already-preprocessed int32
// image held in an ImagePreprocessorBuffer (the same buffer AzIntEngineGPU/ROIIntegrationGPU
// consume): masked/bad pixels are INT32_MIN and saturated pixels INT32_MAX, so bad-pixel identity
// is owned by the preprocessor and a pixel is valid iff v != INT32_MIN && v != INT32_MAX.
//
// The integrator is selected by BraggIntegrationSettings::Integrator:
//   BoxSum          -> BraggIntegrate2D equivalent (rough disk sum minus ring-mean background)
//   ProfileGaussian -> per-reflection measured-width Gaussian profile fit (the default)
//   ProfileEmpirical-> per-shell learned empirical profile fit
// The box sum is also the seed pass (Pass A) of the two profile modes, so it always runs.
//
// This is the Bragg integrator used by the pipeline (bound in MXAnalysisWithoutFPGA: the GPU
// engine when a device is present, otherwise the CPU engine). It takes a preprocessed image +
// the predicted reflections and returns the vector<Reflection> (I, sigma, bkg, partiality, ...)
// that the downstream scaling/merge consumes unchanged.
// =============================================================================

#include <cmath>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <vector>

#include "../../common/BraggIntegrationSettings.h"
#include "../../common/DiffractionExperiment.h"
#include "../../common/DiffractionGeometry.h"
#include "../../common/Reflection.h"
#include "../image_preprocessing/ImagePreprocessorBuffer.h"

namespace bragg_engine {
// Shared with both engines so the CPU and GPU paths stay numerically aligned.
constexpr int    N_SHELL = 6;                 // resolution shells for per-shell profile learning
constexpr double STRONG_I_OVER_SIGMA = 5.0;   // strong-spot threshold that seeds the profile
constexpr int    MIN_STRONG_PER_SHELL = 30;   // below this a shell falls back to the global profile
constexpr double C_CAPTURE = 2.5;             // weak-spot radial capture term (coefficient of tan^2(2theta))
// Lower bound on the background term of the Kabsch fit weights (v = max(bkg, floor) + signal). It
// guards the background ESTIMATE, not the detector: the r2..r3 ring mean of a high-angle reflection
// can come out exactly zero, and v = 0 makes the weights P^2/v diverge. A ring of n pixels cannot
// resolve a background below ~1/n (0.005..0.02 for the default r2=6/r3=10 stencil), so that is the
// scale the floor has to work at. Anything larger over-regularizes: the floor multiplies the reported
// variance by floor/bkg for every pixel below it, so the previous 1/12 inflated sigma by 1.3x at
// 0.05 ct/px and 1.7x at 0.03 - exactly where the weakest high-resolution data live. Digitisation
// noise, where a detector has it, is additive on top of the background and does not belong here.
constexpr double PIXEL_VARIANCE_FLOOR = 0.01;
// The plug-in signal term of the fit weights may lower the per-pixel variance as well as raise it,
// but not below this fraction of the background. Half-wave rectifying it (max(0, I)) instead makes
// the weights - and so the reported 1/den - respond only to upward fluctuations of a noisy intensity
// estimate, which adds ~0.4*sigma*sum(P^3)/sum(P^2)^2 to every sigma whatever the count rate.
constexpr double WEIGHT_VARIANCE_MIN_FRACTION = 0.5;

// Guard against profile-fit runaways: on a weak / near-zero reflection the reweighted Kabsch iteration
// has no real peak to lock onto and can manufacture intensity the box sum never sees. Fall back to the
// summation (box-sum) intensity when the profile result disagrees with the summation seed by more than
// this many box-sum sigmas (a real fit agrees within counting noise, so the margin is generous).
constexpr double PROFILE_SUMMATION_MAX_NSIGMA = 10.0;
}   // namespace bragg_engine

// One reflection's extracted intensity, produced by the derived engine and turned into a
// Reflection by Finalize() (which owns the polarization correction and scale bookkeeping).
struct BraggFitResult {
    float I = 0.0f;
    float sigma = NAN;
    float bkg = 0.0f;
    float observed_x = 0.0f;   // intensity-weighted centroid (BoxSum mode only)
    float observed_y = 0.0f;
    // Variance of everything in `sigma` that is NOT the reflection's own Poisson signal (the
    // background and the error of its estimate). The merge rebuilds each partial's variance at the
    // pooled intensity and needs this term; back-deriving it as sigma^2 - I only works while that
    // identity holds exactly, which it does not for a profile fit.
    float var_bkg = 0.0f;
    bool  ok = false;
    bool  has_observed = false;
};

class BraggIntegrationEngine {
protected:
    // --- fixed configuration extracted from the experiment (see ProfileIntegrate2D) ---
    IntegratorMode mode;
    bool empirical;              // ProfileEmpirical (vs ProfileGaussian)

    size_t xpixel, ypixel, npixel;

    float r1_sq;
    float r2, r2_sq;
    float r3, r3_sq;
    int   R, G, GG;              // profile-grid half-size, edge (2R+1) and area (G*G)

    double bw_sigma;             // bandwidth sigma [dimensionless, * Rpx -> px]
    float  bkg_clip_nsigma;      // high-outlier background sigma-clip multiplier (0 = no clip)
    bool   use_ellipse;          // radially elongate the per-reflection Gaussian

    double c_radial;             // radial variance coefficient of tan^2(2theta): parallax + capture
    double F_px;                 // detector distance expressed in pixels
    float  beam_x, beam_y;
    // Effective symmetric trimmed-mean background fraction (BraggIntegrationSettings), used only when
    // the high-side clip is switched off - the two are alternatives. 0 = plain ring mean. Read by both
    // the CPU and GPU engines.
    float  bkg_trim = 0.0f;

    // --- radial background curvature correction (BraggIntegrationSettings) ---
    // The disk and the annulus are concentric, so any background LINEAR in position cancels between
    // them; what survives is the curvature of the radial background. Every reflection uses the same
    // stencil, so mean_annulus(B) - mean_disk(B) of a radial B is a FIXED kernel over radial offset:
    //     bkg_error = sum_k k_diff[k] * B(r0 + k - k_off)
    // That is one short dot product per reflection and reads no pixels. Built in the constructor.
    // The kernel is built whatever the setting says, so bkg_radial can be flipped between images at
    // no cost - which is what the auto mode does, applying the correction only to the images whose
    // background really is a smooth function of radius (see BackgroundRadial below).
    bool  bkg_radial = false;
    bool  bkg_radial_auto = false;   // settings left it unset: decide per image from the ice score
    int   k_off = 0;                 // index of offset 0 in k_diff
    std::vector<float> k_diff;       // annulus-minus-disk weight per integer radial offset

    DiffractionGeometry geom;    // kept for the per-reflection polarization correction
    std::optional<float> polarization;

    // Assemble output reflections from the per-reflection fit results (polarization + scale corr).
    std::vector<Reflection> Finalize(const std::vector<Reflection> &predicted, size_t npredicted,
                                     const std::vector<BraggFitResult> &results,
                                     int64_t image_number) const;

public:
    explicit BraggIntegrationEngine(const DiffractionExperiment &experiment);
    virtual ~BraggIntegrationEngine() = default;

    // predicted[0..npredicted) are the reflections to extract; image is the preprocessed int32
    // frame (image.size() == npixel). Returns only the observed reflections.
    virtual std::vector<Reflection> Run(const ImagePreprocessorBuffer &image,
                                        const std::vector<Reflection> &predicted, size_t npredicted,
                                        int64_t image_number) = 0;

    // Turn the radial background correction on or off for the images that follow. The caller owns
    // the decision; in the auto mode the analysis sets it per image from that image's ice score.
    void BackgroundRadial(bool on) { bkg_radial = on; }
    [[nodiscard]] bool IsBackgroundRadialAuto() const { return bkg_radial_auto; }
};