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 <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"
#include "../rugnux/StillsLiveMerge.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.
// With a cell given, the pre-pass runs again on all the frames held each time they have doubled, until the geometry
// fit has its frames, the collection ends or the sample reaches PREPASS_MAX_IMAGES or half of `max_held`; the
// analysis starts after that and catches up from the frames held. The bound keeps the last run short and leaves
// room in the queue for the frames arriving meanwhile, which are otherwise not written. The geometry is fixed when the pre-pass ends: every frame is analysed and merged at it. Where the pre-pass could
// not refine it (too few of the frames it had indexed: a small image buffer, a serial hit rate, a chip whose first
// frames are empty; or no cell given), the header's distance stays for the whole run, with a warning in the log
// and the status.
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 int64_t PREPASS_MAX_IMAGES = 4 * PREPASS_IMAGES;
    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_last = false;   // what is held is all there will be
    bool prepass_over = false;
    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);

    const int64_t max_held;
    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;

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