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 <condition_variable>
#include <exception>
#include <functional>
#include <future>
#include <latch>
#include <map>
#include <memory>
#include <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"

// 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.
class StillsGPUAnalysis {
public:
    constexpr static int64_t PREPASS_IMAGES = 100;

    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);
    ~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.
    [[nodiscard]] std::optional<UnitCell> GetConsensusUnitCell() const;
    [[nodiscard]] float GetBeamX_pxl() const;
    [[nodiscard]] float GetBeamY_pxl() const;
    [[nodiscard]] float GetDetectorDistance_mm() 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;

    std::mutex live_mutex;
    RugnuxLiveStills live;
    std::chrono::steady_clock::time_point ready_time;
    void Publish();

    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;

    void StartPrePass();
    void PrePassThread();
    void AnalysisThread();
};