Jungfraujoch/image_analysis/beam_stop/ShadowFinder.h

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

#pragma once

#include <cstdint>
#include <future>
#include <memory>
#include <mutex>
#include <optional>
#include <vector>

#include "../../common/CompressedImage.h"
#include "../../common/DiffractionExperiment.h"
#include "../../common/DiffractionGeometry.h"
#include "../../common/JFJochMessages.h"
#include "../../common/PixelMask.h"
#ifdef JFJOCH_USE_CUDA
#include "ShadowAccumulatorGPU.h"
#endif

// Finds the beam-stop shadow - the central disk and the holder arm - from a set of images,
// mirroring the accumulate-then-finalize shape of DarkMaskAnalysis: feed frames with
// AddImage(), then read the mask once with GetMask(). The mask is in converted geometry
// and is 1 where the beam stop shadows the detector.
//
// The shadow is a place where the background is missing, so it is found by comparing each
// pixel's mean against the typical background at the same radius - the median over its ring,
// taken over the pixels not already known to be shadowed. That comparison holds wherever the
// ring still has unshadowed pixels to measure. Where it does not - a ring lying wholly inside
// the stop - there is nothing to compare against, and such a ring is shadow in its entirety.
//
// The background belongs to the beam and the shadow to the stop, and the two are not concentric:
// the stop sits off the beam by a sizeable fraction of its own radius. Only the per-ring
// comparison is used, so nothing here assumes they share a centre.
//
// Frames are chosen by the caller; the detection needs enough of them that the background
// is counted rather than guessed (see MIN_EXPECTED_COUNTS in the .cpp).
// Thread-safe: workers call AddImage concurrently, each naming a shard of its own (see
// SetShardCount) - so no two threads touch the same accumulator and nothing is locked while
// an image is added. The shards are summed when the projection is read.
class ShadowFinder {
    mutable std::mutex m;

    const int   width;
    const int   height;
    float       beam_x;
    float       beam_y;

    // What the background owes to the source rather than to the hardware. The scattered background
    // is not flat around a ring: a polarized source suppresses it in its own plane by a factor that
    // reaches three at the 2 theta a short detector distance puts in a corner - several times the
    // dip this class is looking for - so the comparison divides it out before it compares. Of the
    // corrections a ring carries this is the only one that varies along it; solid angle, detector
    // and air absorption are all functions of 2 theta alone and the ring's own median absorbs them.
    // The geometry is kept whole rather than reduced to a distance and a pixel size because the
    // azimuth is the whole point: detector tilt, a quarter-turned image and an in-plane rotation
    // all move the polarization plane across the stored image, and the geometry already knows where
    // it lies. Only the centre is replaced, by the one the caller measured.
    const DiffractionGeometry geometry;
    const std::optional<float> polarization;   // unset leaves the background as it was measured

    std::vector<uint32_t> pixel_mask;   // pixels already masked carry no background to test

    // Per-pixel projection over the frames added so far (converted geometry). One set per shard:
    // the sums and counts are integers, so summing the shards is exact and the result does not
    // depend on how the frames were spread over them.
    struct Projection {
        std::vector<int64_t>  max_value;
        std::vector<int64_t>  sum_value;
        std::vector<uint32_t> valid_count;
        uint32_t frames = 0;
    };
    std::vector<Projection> shards;

#ifdef JFJOCH_USE_CUDA
    // Present when a GPU is available. Frames it can decode are accumulated there instead of on the
    // host - only the compressed chunk crosses PCIe - and its projection is folded in with the
    // shards when the mask is read. Frames it cannot take (anything but bitshuffle+LZ4) still go to
    // a host shard, so a run mixing compressions is handled without a second code path.
    // Built on a thread of its own: it allocates and clears several hundred megabytes of device
    // memory, and cudaMalloc synchronises the whole device, so doing it in the constructor would
    // stall the caller before it has read its first frame. The first AddImage waits for it, by
    // which time the reads have been running for a while.
    mutable std::future<std::unique_ptr<ShadowAccumulatorGPU>> gpu_pending;
    mutable std::unique_ptr<ShadowAccumulatorGPU> gpu;
    mutable std::mutex gpu_mutex;

    // The accumulator once its construction has finished, or null if there is none.
    [[nodiscard]] ShadowAccumulatorGPU *Gpu() const;
#endif

    template<class T> void Add(const T *ptr, Projection &p);

    // Sum the shards into one projection. max_value is only taken from a shard that actually
    // counted the pixel - a shard that never saw it holds 0, which would beat a genuinely
    // negative maximum.
    [[nodiscard]] Projection Reduce() const;

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

    // The centre the rings are drawn about. It starts as the file's, which is the only one there
    // is when the finder is built; a caller that has measured one replaces it before reading the
    // mask. A centre far from the truth draws the rings across the background's own radial
    // fall-off instead of along it, and the comparison then describes the fall-off rather than the
    // hardware. The projection is not centred on anything, so this may be set after the frames.
    void BeamCenter(float x, float y);

    // Give each worker a shard to accumulate into. Must be called before the first AddImage,
    // and costs 20 bytes per pixel per shard.
    void SetShardCount(size_t n);

    // Accumulate one full converted-geometry image into shard `shard`. Gap / masked pixels
    // (the pixel type's sentinel extreme) are skipped. `buffer` is scratch space for
    // decompression, reused across the calls of one worker.
    void AddImage(const DataMessage &data, std::vector<uint8_t> &buffer, size_t shard = 0);

    // Compute the shadow mask (1 = shadow, 0 = keep), of the converted pixel count.
    // Recomputed from the accumulators on each call - meant to be called once at the end.
    // nthreads = 0 asks for all hardware threads. The per-pixel passes over a 16M-pixel detector
    // dominate this, and they are all exactly parallel.
    [[nodiscard]] std::vector<uint32_t> GetMask(size_t nthreads = 0) const;

    // Mean counts per pixel over the frames added, NAN where nothing was counted. This is the
    // projection GetMask() tests, so anything else that wants the background before indexing
    // gets it without reading the frames a second time.
    [[nodiscard]] std::vector<float> GetMeanProjection() const;

    [[nodiscard]] uint32_t GetFrameCount() const;
};