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

#include "RugnuxLive.h"

#include <future>

#include "RugnuxDefaults.h"
#include "../common/CUDAWrapper.h"
#include "../image_analysis/indexing/IndexerThreadPool.h"
#ifdef JFJOCH_USE_CUDA
#include "../image_analysis/geom_refinement/BeamCenterFFTGPU.h"
#endif

RugnuxLiveSetup MakeRugnuxLiveSetup(const JFJochReaderDataset &dataset, int nthreads,
                                    std::shared_ptr<const RugnuxModel> model) {
    RugnuxLiveSetup s;

    // The experiment as rugnux_cli.cpp sets it up for a rotation sweep with no options given.
    DiffractionExperiment &x = s.experiment;
    x = dataset.experiment;
    ApplyRugnuxExperimentDefaults(x);
    x.NumTriggers(1);
    ClearStoredCrystal(x);
    x.MaxSpotCount(RUGNUX_MAX_SPOT_COUNT);

    IndexingSettings indexing;
    indexing.Algorithm(IndexingAlgorithmEnum::Auto);
    indexing.RotationIndexing(true);
    indexing.GeomRefinementAlgorithm(GeomRefinementAlgorithmEnum::BeamCenter);
    indexing.IndexIceRings(false);
    x.ImportIndexingSettings(indexing);
    x.DetectIceRings(RugnuxDetectIceRings(dataset.file_detect_ice_rings, true));

    ScalingSettings scaling = RugnuxDefaultScalingSettings(true);
    scaling.ScaleFulls(true);
    scaling.RelativeBDegrees(0.0);
    scaling.MergeFriedel(true);
    scaling.MinPartiality(0.02);
    scaling.MinCCForImage(0.0);
    x.ImportScalingSettings(scaling);

    ProcessConfig &c = s.config;
    c.mode = ProcessMode::FullAnalysis;
    c.nthreads = nthreads;
    // As a bare `rugnux` with no -N: the per-image loop keeps to what a card can take (ImagePass).
    c.nthreads_auto = true;
    c.export_unmerged = false;
    c.spot_finding = RugnuxDefaultSpotFinding(true);
    c.rotation_indexing = true;
    c.two_pass_rotation = true;
    c.rotation_postrefine_geometry = true;
    c.rotation_indexing_image_count = RUGNUX_ROTATION_INDEXING_IMAGES;
    c.detect_beam_stop = RUGNUX_DETECT_BEAM_STOP_FRAMES;
    c.beam_center_check = true;
    c.measure_spot_budget = true;
    c.run_scaling = true;
    c.model = std::move(model);
    c.live = true;
    c.fast_live = true;
    // Decided late, on the whole sweep, as it re-checks its radius there - which fast_live does not.
    c.adaptive_integration_radius = false;
    return s;
}

void WarmUpRugnuxLiveGPU(const RugnuxLiveSetup &setup) {
#ifdef JFJOCH_USE_CUDA
    // As Rugnux sets them for its first-pass pool (the measured centre's indexing and pass 2 use it as
    // well) and its short-axis pool. Should they part, the run builds its own, as without this.
    IndexingSettings first_pass = setup.experiment.GetIndexingSettings();
    first_pass.RefineThreads(std::clamp(setup.config.nthreads / 2, 1, MAX_REFINE_THREADS));
    IndexingSettings short_axis = first_pass;
    short_axis.FFT_MinUnitCell_A(IndexingSettings::fft_min_unit_cell_limit_A);

    // Each pool has two workers, on consecutive cards.
    const int gpus = get_gpu_count();
    const int per_card = gpus == 1 ? 2 : 1;
    std::vector<std::future<void>> cards;
    for (int32_t i = 0; i < gpus; i++)
        cards.push_back(std::async(std::launch::async, [&, i] {
            ScopedGpuPin pin(i);
            KeepFFTIndexers(first_pass, per_card);
            KeepFFTIndexers(short_axis, per_card);
            // Sizing the capture's work area makes its two plans (CudaFFTPlan::WorkBytes).
            BeamCenterFFTGPU::DeviceMemoryNeeded(setup.experiment.GetXPixelsNumConv(),
                                                 setup.experiment.GetYPixelsNumConv());
        }));
    for (auto &c: cards)
        c.get();
#endif
}