Jungfraujoch/receiver/StillsGPUAnalysis.h

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

#pragma once

#include <array>
#include <atomic>
#include <condition_variable>
#include <exception>
#include <functional>
#include <future>
#include <latch>
#include <map>
#include <memory>
#include <mutex>
#include <shared_mutex>
#include <vector>

#include "StillsPrePass.h"
#include "RugnuxLiveStatus.h"
#include "../common/DiffractionExperiment.h"
#include "../common/JFJochException.h"
#include "../common/AzimuthalIntegrationMapping.h"
#include "../common/AzimuthalIntegrationProfile.h"
#include "../common/Logger.h"
#include "../common/ThreadSafeFIFO.h"
#include "../image_analysis/IndexAndRefine.h"
#include "../rugnux/StillsLiveMerge.h"
#include "../rugnux/StillsRunningGeometry.h"

// Serial frames analysed on the GPU as the DECTRIS ones are (MXAnalysisWithoutFPGA), after a pre-pass on
// the first PREPASS_IMAGES (RunStillsPrePass) has decided the beam centre and the mask they are analysed
// with. Every frame given before the pre-pass is done is held; `ready` is called once it is, with the
// geometry and mask the frames are then analysed at, before any of them is analysed; `sink` gets each
// frame once analysed, from one of the analysis threads. Where `status` is given, the run is reported there
// (pre-pass, its decisions, the analysis).
// At most `max_held` frames wait for analysis: Hold refuses one more, so the slots held stay bounded.
// Where the setup after the pre-pass fails (`ready` included), `failed` is called at once, the frames are
// dropped unanalysed and Finish throws the error. An analysis thread that fails passes its frames on
// unanalysed, and the run is reported as failed.
// The indexed frames are merged as `rugnux --mode scale` merges them (StillsLiveMerge), beside the analysis:
// a provisional merge at each of CHECKPOINTS of the images, on MERGE_THREADS threads, reported as a step
// of the status, and the merge of the whole collection once the last frame is analysed, on every thread,
// before Finish returns (GetMerged). The cell is the run's consensus at the first merge, as rugnux merges in
// it (the collection's cell picks the frames merged); the space group the collection's, else P1.
// merge = false leaves it out.
// Where the pre-pass refined the centre and distance on fewer frames than StillsRunningGeometry bundles, or
// could not refine them (too few of the frames it had indexed: a small image buffer, a serial hit rate, a chip
// whose first frames are empty), and a cell is given, they are refined again while the frames are analysed
// (StillsRunningGeometry, from the spots and lattices of the frames analysed, no image held), and the analysis
// switches to that geometry as each of GEOMETRY_ROUNDS fits is made. The frames analysed before the last
// switch are not merged (their reflections were predicted at an earlier geometry; the images are gone); where
// no fit can be made they are merged as analysed.
class StillsGPUAnalysis {
public:
    // Enough frames for the distance to be refined at a serial hit rate: of 100 frames of a collection indexed
    // at 18%, 7 indexed at a distance 2.6% off and the geometry refinement did not run; of 500, 34 did and
    // refined it to 0.4%. They wait in the image buffer, which has to hold them in half its slots.
    constexpr static int64_t PREPASS_IMAGES = 500;
    constexpr static std::array<double, 2> CHECKPOINTS = {0.5, 0.75};
    constexpr static int MERGE_THREADS = 8;

    struct Setup {
        DiffractionExperiment experiment;
        PixelMask mask;
        std::unique_ptr<AzimuthalIntegrationMapping> mapping;
    };
    using Ready = std::function<void(const Setup &)>;
    using Sink = std::function<void(DataMessage &, const AzimuthalIntegrationProfile &)>;
    using Settings = std::function<SpotFindingSettings()>;
    using Failed = std::function<void(const JFJochException &)>;

    // `start`: the collection's start message, as it would be sent without the pre-pass.
    StillsGPUAnalysis(const DiffractionExperiment &experiment, const PixelMask &mask, const StartMessage &start,
                      IndexerThreadPool *indexing_thread_pool, int64_t nthreads, int64_t max_held, Logger &logger,
                      Ready ready, Sink sink, Settings settings, Failed failed = {},
                      RugnuxLiveStatus *status = nullptr, bool merge = true);
    ~StillsGPUAnalysis();

    // False where `max_held` frames already wait: the frame is not taken.
    bool Hold(const std::shared_ptr<HeldImage> &image);
    // No frame comes after this: the pre-pass runs on what came, if it has not, and every frame is
    // analysed before it returns. Throws where the setup after the pre-pass failed.
    void Finish();
    // After Finish, once the end message is out and the writer is done: the status gets the latency from the last
    // image and then the run's outcome. Without it the outcome is reported when the analysis goes, with no latency.
    void Written();

    // After Finish.
    [[nodiscard]] std::optional<UnitCell> GetConsensusUnitCell() const;
    [[nodiscard]] float GetBeamX_pxl() const;
    [[nodiscard]] float GetBeamY_pxl() const;
    [[nodiscard]] float GetDetectorDistance_mm() const;
    // The merge of the collection, where one was made.
    [[nodiscard]] const std::optional<StillsMergeOutput> &GetMerged() const;
    [[nodiscard]] const DiffractionExperiment &GetMergeExperiment() const;

private:
    const DiffractionExperiment &experiment;
    const PixelMask &mask;
    const StartMessage start;
    IndexerThreadPool *indexing_thread_pool;
    Logger &logger;
    Ready ready;
    Sink sink;
    Settings settings;
    Failed setup_failed;
    RugnuxLiveStatus *status;
    const bool merge_frames;

    std::mutex live_mutex;
    RugnuxLiveStills live;
    std::chrono::steady_clock::time_point ready_time;
    void Publish();
    // Kept by Finish for Written
    std::optional<std::chrono::steady_clock::time_point> last_image;
    std::optional<std::pair<RugnuxLiveState, std::string>> outcome;

    Setup setup;
    std::unique_ptr<IndexAndRefine> indexer;

    std::mutex prepass_mutex;
    std::condition_variable prepass_cv;
    bool prepass_go = false;
    int64_t images_held = 0;
    std::map<int64_t, std::shared_ptr<HeldImage>> prepass_images;
    std::latch prepass_done{1};
    std::exception_ptr setup_error; // set before prepass_done

    std::mutex error_mutex;
    std::string analysis_error; // the first, where an analysis thread failed
    void AnalysisError(const std::string &text);

    ThreadSafeFIFO<std::shared_ptr<HeldImage>> queue;
    std::future<void> prepass_future;
    std::vector<std::future<void>> analysis_futures;
    bool finished = false;

    // The frames analysed, by number, for the merge (-1 after the last); it takes their outcomes from the indexer.
    ThreadSafeFIFO<int64_t> merge_queue;
    std::unique_ptr<StillsLiveMerge> merge;   // after the indexer: its reflections are in the indexer's arena
    DiffractionExperiment merge_experiment;
    std::optional<StillsMergeOutput> merged;
    std::future<void> merge_future;

    // The geometry refined while the frames are analysed (see above), in GEOMETRY_ROUNDS rounds, each on the
    // frames analysed at the last one's geometry (the spots of a frame are indexed at the geometry it was
    // analysed at, and a start 2.6% off left the first fit 0.6% off). The analysis takes geometry_mutex shared
    // for each frame; a switch takes it exclusive and raises geometry_epoch, which the analysis threads see to
    // rebuild their analysis. geometry_state is PENDING until the last round, then SWITCHED or KEPT.
    // frame_epoch: per frame number, 1 + the epoch it was analysed at while PENDING (0: after); the merge
    // leaves out those analysed before the last switch.
    constexpr static int GEOMETRY_ROUNDS = 2;
    enum GeometryState : int { GEOMETRY_PENDING = 0, GEOMETRY_SWITCHED = 1, GEOMETRY_KEPT = 2 };
    std::unique_ptr<StillsRunningGeometry> running_geometry;
    int geometry_round = 0;
    std::shared_mutex geometry_mutex;
    // A switch waiting for geometry_mutex closes this gate to new frames, or the frames that keep taking the
    // mutex shared would hold it off to the end of the collection.
    std::mutex gate_mutex;
    std::condition_variable gate_cv;
    bool switch_waiting = false;
    std::atomic<int> geometry_epoch{0};
    std::atomic<int> geometry_state{GEOMETRY_KEPT};
    std::vector<uint8_t> frame_epoch;
    std::vector<std::unique_ptr<AzimuthalIntegrationMapping>> retired_mappings;   // the analyses before the switch read them
    void SwitchGeometry();

    void StartPrePass();
    void PrePassThread();
    void AnalysisThread();
    void MergeThread();
    void PublishMerge(const StillsMergeOutput &out, const std::string &title);
};