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 <vector>

#include "../../common/CompressedImage.h"
#include "../../common/DiffractionExperiment.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;
    const float beam_x;
    const float beam_y;

    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);

    // 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;
};