Jungfraujoch/receiver/StillsGPUAnalysis.cpp

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

#include "StillsGPUAnalysis.h"

#include <thread>

#include "../common/CUDAWrapper.h"
#include "../common/ThreadAffinity.h"
#include "../image_analysis/MXAnalysisWithoutFPGA.h"
#include "../rugnux/RugnuxDefaults.h"

StillsGPUAnalysis::StillsGPUAnalysis(const DiffractionExperiment &in_experiment, const PixelMask &in_mask,
                                     const StartMessage &in_start, IndexerThreadPool *in_indexing_thread_pool,
                                     int64_t nthreads, int64_t in_max_held, Logger &in_logger, Ready in_ready,
                                     Sink in_sink, Settings in_settings, Failed in_failed,
                                     RugnuxLiveStatus *in_status, bool in_merge)
    : experiment(in_experiment), mask(in_mask), start(in_start), indexing_thread_pool(in_indexing_thread_pool),
      logger(in_logger), ready(std::move(in_ready)), sink(std::move(in_sink)), settings(std::move(in_settings)),
      setup_failed(std::move(in_failed)), status(in_status), merge_frames(in_merge),
      max_held(std::max<int64_t>(1, in_max_held)), queue(max_held) {
    live.images_needed = std::min(PREPASS_IMAGES, experiment.GetImageNum());
    live.file_beam_x_pxl = experiment.GetBeamX_pxl();
    live.file_beam_y_pxl = experiment.GetBeamY_pxl();
    live.file_distance_mm = experiment.GetDetectorDistance_mm();
    if (status) {
        status->Start(experiment.GetImageNum(),
                      static_cast<double>(experiment.GetImageNum())
                      * std::chrono::duration<double>(experiment.GetImageTime()).count(),
                      std::chrono::steady_clock::now());
        status->Step("Pre-pass");
        Publish();
    }
    prepass_future = std::async(std::launch::async, &StillsGPUAnalysis::PrePassThread, this);
    for (int64_t i = 0; i < nthreads; i++)
        analysis_futures.emplace_back(std::async(std::launch::async, &StillsGPUAnalysis::AnalysisThread, this));
    logger.Info("GPU analysis: {} threads, at least the first {} images held for the pre-pass", nthreads, PREPASS_IMAGES);
}

StillsGPUAnalysis::~StillsGPUAnalysis() {
    try {
        if (!finished)
            Finish();
    } catch (const std::exception &e) {
        logger.Error("GPU analysis: {}", e.what());
    }
    if (status && outcome)
        status->Finish(outcome->first, outcome->second);
}

void StillsGPUAnalysis::AnalysisError(const std::string &text) {
    logger.Error("GPU analysis: {}", text);
    std::unique_lock ul(error_mutex);
    if (analysis_error.empty())
        analysis_error = text;
}

void StillsGPUAnalysis::Publish() {
    std::unique_lock ul(live_mutex);
    if (live.prepass == RugnuxLivePrePass::Done || live.prepass == RugnuxLivePrePass::Failed) {
        const double t = std::chrono::duration<double>(std::chrono::steady_clock::now() - ready_time).count();
        if (t > 0 && live.images_analysed > 0)
            live.analysis_rate_hz = static_cast<double>(live.images_analysed) / t;
    }
    status->Stills(live);
}

bool StillsGPUAnalysis::Hold(const std::shared_ptr<HeldImage> &image) {
    if (status)
        status->ImageReceived(std::chrono::steady_clock::now());
    if (!queue.Put(image)) {
        // As many frames wait as the buffer is to hold: the pre-pass runs on what it has, for the last time.
        StartPrePass();
        return false;
    }
    {
        std::unique_lock ul(prepass_mutex);
        if (!prepass_over) {
            prepass_images[image->message.number] = image;
            std::unique_lock ul_live(live_mutex);
            live.images_held = static_cast<int64_t>(prepass_images.size());
        }
    }
    prepass_cv.notify_all();
    if (status)
        Publish();
    return true;
}

void StillsGPUAnalysis::StartPrePass() {
    {
        std::unique_lock ul(prepass_mutex);
        prepass_last = true;
    }
    prepass_cv.notify_all();
}

void StillsGPUAnalysis::Finish() {
    finished = true;
    StartPrePass();
    prepass_future.get();
    for (size_t i = 0; i < analysis_futures.size(); i++)
        queue.PutBlocking(nullptr);
    for (auto &f: analysis_futures)
        f.get();
    if (status) {
        last_image = status->LastImage();
        Publish();
    }
    if (merge_future.valid()) {
        if (status)
            status->Step("Merge");
        merge_queue.Put(-1);
        merge_future.get();
        if (merged && status) {
            PublishMerge(*merged, "the collection");
            status->Headline(merge_experiment.GetSpaceGroupOrP1().xhm(), merged->merge.statistics.overall.d_min,
                             std::nullopt);
        }
    }
    if (status) {
        if (setup_error)
            status->Finish(RugnuxLiveState::Failed, "GPU analysis setup failed");
        else {
            // Reported by Written, once everything is written - the end message and the writer go on meanwhile.
            status->Step("Writing");
            std::unique_lock ul(error_mutex);
            if (analysis_error.empty())
                outcome.emplace(RugnuxLiveState::Done, "");
            else
                outcome.emplace(RugnuxLiveState::Failed, analysis_error);
        }
    }
    if (setup_error)
        std::rethrow_exception(setup_error);
}

void StillsGPUAnalysis::Written() {
    if (!status || !outcome)
        return;
    if (last_image)
        status->Latency(std::chrono::duration<double>(std::chrono::steady_clock::now() - last_image.value()).count());
    status->Finish(outcome->first, outcome->second);
    outcome.reset();
}

void StillsGPUAnalysis::PrePassThread() {
    const auto masked = [](const PixelMask &m) {
        return std::count_if(m.GetMask().begin(), m.GetMask().end(), [](uint32_t v) { return v != 0; });
    };
    setup.experiment = experiment;
    setup.mask = mask;
    const int64_t all_images = experiment.GetImageNum();
    std::optional<StillsPrePassResult> result;
    std::chrono::steady_clock::time_point prepass_start = std::chrono::steady_clock::now();
    // At a low hit rate the first frames may not be enough to fit the geometry on, or may not hold a crystal: with
    // a cell to fit, the pre-pass runs again on all that is held once that has doubled, until the fit has its
    // frames, or the collection ended, or the sample is as large as it may grow (the last try): bounded so that
    // the last run is short and the other half of the queue takes the frames arriving meanwhile.
    const int64_t max_sample = std::min(PREPASS_MAX_IMAGES, max_held / 2);
    for (int64_t needed = std::min(PREPASS_IMAGES, all_images); ;) {
        std::map<int64_t, std::shared_ptr<HeldImage>> sample;
        bool last;
        {
            std::unique_lock ul(prepass_mutex);
            prepass_cv.wait(ul, [&] {
                return prepass_last || static_cast<int64_t>(prepass_images.size()) >= needed;
            });
            // Frames that arrived during the previous run do not make this one larger than the bound.
            for (const auto &[n, image]: prepass_images) {
                if (static_cast<int64_t>(sample.size()) >= std::max(needed, max_sample))
                    break;
                sample.emplace(n, image);
            }
            last = prepass_last || static_cast<int64_t>(sample.size()) >= std::min(all_images, max_sample);
        }
        if (sample.empty())
            break;
        {
            std::unique_lock ul(live_mutex);
            live.prepass = RugnuxLivePrePass::Running;
        }
        if (status)
            Publish();
        try {
            const auto t0 = std::chrono::steady_clock::now();
            auto r = RunStillsPrePass(start, sample, static_cast<int>(std::thread::hardware_concurrency()));
            logger.Info("Pre-pass on {} images ({:.1f} s): beam centre ({:.2f}, {:.2f}) -> ({:.2f}, {:.2f}) px, "
                        "distance {:.2f} -> {:.2f} mm on {} frames, masked pixels {} -> {}", r.images,
                        std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count(),
                        experiment.GetBeamX_pxl(), experiment.GetBeamY_pxl(), r.beam_x_pxl, r.beam_y_pxl,
                        experiment.GetDetectorDistance_mm(), r.distance_mm, r.geometry_frames,
                        masked(mask), masked(r.mask));
            result = std::move(r);
        } catch (const std::exception &e) {
            logger.Warning("Pre-pass on {} images failed - {}", sample.size(), e.what());
        }
        if (last || !experiment.GetUnitCell()
            || (result && result->geometry_frames >= RUGNUX_REFINE_GEOMETRY_FRAMES))
            break;
        needed = std::min({2 * static_cast<int64_t>(sample.size()), all_images, max_sample});
        std::unique_lock ul(live_mutex);
        live.prepass = RugnuxLivePrePass::Holding;
        live.images_needed = needed;
    }
    {
        // The frames are analysed from here on: their slots go back to the buffer once they are sent.
        std::unique_lock ul(prepass_mutex);
        prepass_over = true;
        prepass_images.clear();
    }
    const bool failed = !result;
    int geometry_frames = 0;
    if (result) {
        const auto &r = *result;
        setup.experiment.BeamX_pxl(r.beam_x_pxl).BeamY_pxl(r.beam_y_pxl).DetectorDistance_mm(r.distance_mm);
        setup.mask = r.mask;
        geometry_frames = r.geometry_frames;
        const auto &m = r.mask.GetMask();
        std::unique_lock ul(live_mutex);
        live.beam_x_pxl = r.beam_x_pxl;
        live.beam_y_pxl = r.beam_y_pxl;
        live.distance_mm = r.distance_mm;
        live.shadow_pixels = std::count_if(m.begin(), m.end(), [](uint32_t v) {
            return (v & (1u << PixelMask::BeamStopPixelBit)) != 0;
        });
        live.bad_pixels = std::count_if(m.begin(), m.end(), [](uint32_t v) {
            return (v & (1u << PixelMask::HotPixelBit)) != 0;
        });
    } else
        logger.Warning("Pre-pass failed, the images are analysed with the stated beam centre and mask");
    try {
        setup.mapping = std::make_unique<AzimuthalIntegrationMapping>(setup.experiment, setup.mask);
        // Retained for the merge, which takes each image's reflections as soon as it is analysed.
        indexer = std::make_unique<IndexAndRefine>(setup.experiment, indexing_thread_pool, merge_frames, true);
        if (geometry_frames == 0) {
            const std::string text = fmt::format(
                    "The pre-pass could not refine the geometry (too few frames indexed in what the buffer holds, or no unit cell given): the "
                    "whole run is analysed at the distance in the header, {:.2f} mm, beam centre ({:.2f}, {:.2f}) px",
                    setup.experiment.GetDetectorDistance_mm(), setup.experiment.GetBeamX_pxl(),
                    setup.experiment.GetBeamY_pxl());
            logger.Warning("{}", text);
            if (status)
                status->Note(text);
        }
        if (merge_frames) {
            merge_experiment = setup.experiment;
            merge_experiment.ImportScalingSettings(RugnuxDefaultScalingSettings(false));
            merge = std::make_unique<StillsLiveMerge>(merge_experiment, experiment.GetImageNum(),
                                                      setup.experiment.GetUnitCell(), std::nullopt,
                                                      setup.experiment.IsDetectIceRings(),
                                                      SpotFindingSettings().ice_ring_width_Q_recipA, logger);
            merge_future = std::async(std::launch::async, &StillsGPUAnalysis::MergeThread, this);
        }
        ready(setup);
    } catch (const std::exception &e) {
        // Nothing can be written without the start message: the frames are dropped (AnalysisThread), and
        // the collection is stopped.
        const auto *j = dynamic_cast<const JFJochException *>(&e);
        const JFJochException error = j ? *j : JFJochException(JFJochExceptionCategory::WrongDAQState,
                                                               std::string("GPU analysis setup failed: ") + e.what());
        logger.ErrorException(error);
        setup_error = j ? std::current_exception() : std::make_exception_ptr(error);
        if (setup_failed)
            setup_failed(error);
        prepass_done.count_down();
        return;
    }
    {
        std::unique_lock ul(live_mutex);
        live.prepass = failed ? RugnuxLivePrePass::Failed : RugnuxLivePrePass::Done;
        live.prepass_s = std::chrono::duration<double>(std::chrono::steady_clock::now() - prepass_start).count();
        ready_time = std::chrono::steady_clock::now();
    }
    if (status) {
        status->Step("GPU analysis");
        Publish();
    }
    prepass_done.count_down();
}

void StillsGPUAnalysis::AnalysisThread() {
    try {
        pin_gpu();
    } catch (const std::exception &e) {
        logger.Warning("Error pinning GPU {}", e.what());
    }

    prepass_done.wait();

    if (setup_error) {
        while (queue.GetBlocking()) {}
        return;
    }

    // A thread whose analysis cannot be built still passes its frames on, unanalysed.
    std::unique_ptr<MXAnalysisWithoutFPGA> analysis;
    try {
        analysis = std::make_unique<MXAnalysisWithoutFPGA>(setup.experiment, *setup.mapping, setup.mask,
                                                           *indexer, /*enable_fused_adaptive_gpu=*/true);
    } catch (const std::exception &e) {
        AnalysisError(std::string("analysis thread could not start, its frames are written unanalysed - ")
                      + e.what());
    }

    while (auto held = queue.GetBlocking()) {
        DataMessage &message = held->message;
        std::optional<AzimuthalIntegrationProfile> profile;
        const auto start_time = std::chrono::steady_clock::now();
        profile.emplace(*setup.mapping);
        try {
            if (analysis)
                analysis->Analyze(message, *profile, settings());
        } catch (const JFJochException &e) {
            // The frame is still written, without its analysis.
            logger.ErrorException(e);
        } catch (const std::exception &e) {
            AnalysisError(fmt::format("image {} written unanalysed - {}", message.number, e.what()));
        }
        message.processing_time_s = std::chrono::duration<float>(std::chrono::steady_clock::now()
                                                                  - start_time).count();
        try {
            sink(message, *profile);
            if (merge)
                merge_queue.Put(message.number);
            if (status) {
                int64_t analysed;
                {
                    std::unique_lock ul(live_mutex);
                    analysed = ++live.images_analysed;
                    if (message.indexing_result.value_or(false))
                        ++live.images_indexed;
                }
                status->Progress(analysed, experiment.GetImageNum());
                Publish();
            }
        } catch (const std::exception &e) {
            // The thread goes on with the next frame, so that the queue is drained.
            AnalysisError(fmt::format("image {} lost - {}", message.number, e.what()));
        }
    }
}

std::optional<UnitCell> StillsGPUAnalysis::GetConsensusUnitCell() const {
    return indexer->GetConsensusUnitCell();
}

float StillsGPUAnalysis::GetBeamX_pxl() const {
    return setup.experiment.GetBeamX_pxl();
}

float StillsGPUAnalysis::GetBeamY_pxl() const {
    return setup.experiment.GetBeamY_pxl();
}

float StillsGPUAnalysis::GetDetectorDistance_mm() const {
    return setup.experiment.GetDetectorDistance_mm();
}

const std::optional<StillsMergeOutput> &StillsGPUAnalysis::GetMerged() const {
    return merged;
}

const DiffractionExperiment &StillsGPUAnalysis::GetMergeExperiment() const {
    return merge_experiment;
}

void StillsGPUAnalysis::MergeThread() {
    constexpr int64_t PREPARE_EVERY = 200;   // frames between two Prepare calls after a checkpoint
    const int64_t images = experiment.GetImageNum();
    // The cell merged in, as rugnux merges it: the consensus of the frames indexed by the first merge (the
    // collection's cell only picks the frames that are merged).
    bool cell_set = false;
    const auto have_cell = [&] {
        if (!cell_set)
            if (const auto c = indexer->GetConsensusUnitCell()) {
                merge_experiment.SetUnitCell(*c);
                merge->Cell(*c);
                cell_set = true;
            }
        return cell_set;
    };
    size_t next = 0;
    try {
        const auto add = [&](int64_t n, IntegrationOutcome &&o) {
            merge->Add(n, std::move(o));
            if (next < CHECKPOINTS.size()
                && merge->Added() >= std::llround(CHECKPOINTS[next] * static_cast<double>(images))) {
                if (have_cell())
                    PublishMerge(merge->Checkpoint(MERGE_THREADS),
                                 fmt::format("{:.0f}% of the images (provisional)", 100.0 * CHECKPOINTS[next]));
                ++next;
            } else if (next > 0 && merge->Added() % PREPARE_EVERY == 0)
                merge->Prepare(MERGE_THREADS);
        };
        for (int64_t n = merge_queue.GetBlocking(); n >= 0; n = merge_queue.GetBlocking())
            add(n, indexer->TakeIntegrationOutcome(n));
        if (have_cell()) {
            merged = merge->Final(static_cast<int>(AvailableCpus()));
            logger.Info("Stills merge after the last image: {:.2f} s", merged->seconds);
        }
    } catch (const std::exception &e) {
        logger.Warning("Stills merge failed: {}", e.what());
        merged.reset();
    }
}

void StillsGPUAnalysis::PublishMerge(const StillsMergeOutput &out, const std::string &title) {
    const auto &o = out.merge.statistics.overall;
    const std::string text = fmt::format("Merge of {}: {} crystals, {:.2f} A, CC1/2 {:.3f}, R_split {:.3f}, "
                                         "completeness {:.1f}%, <I/sigma> {:.1f} ({:.1f} s)", title, out.images,
                                         o.d_min, o.cc_half, o.r_split,
                                         o.possible_unique_reflections > 0
                                         ? 100.0 * o.unique_reflections / o.possible_unique_reflections : 0.0,
                                         o.mean_i_over_sigma, out.seconds);
    logger.Info("{}", text);
    if (status)
        status->Note(text);
}