Jungfraujoch/image_analysis/IndexAndRefine.h

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

#pragma once

#include <atomic>
#include <vector>
#include <mutex>
#include <functional>

#include "../common/DiffractionSpot.h"
#include "../common/DiffractionExperiment.h"
#include "../common/AzimuthalIntegrationMapping.h"
#include "../common/AzimuthalIntegrationProfile.h"
#include "../common/Reflection.h"
#include "bragg_prediction/BraggPrediction.h"
#include "indexing/IndexerThreadPool.h"
#include "lattice_search/LatticeSearch.h"
#include "rotation_indexer/RotationIndexer.h"
#include "rotation_indexer/RotationIndexerCounter.h"
#include "scale_merge/ReindexAmbiguity.h"
#include "scale_merge/ScaleOnTheFly.h"
#include "scale_merge/ScalingResult.h"
#include "IntegrationOutcome.h"

// Integrates the predicted reflections off whatever image the caller holds: the preprocessed GPU/CPU
// buffer on the WithoutFPGA path (GPU when available), or the assembled detector image read straight,
// on the CPU, on the forced-CPU FPGA path. Keeps IndexAndRefine independent of the image representation.
using BraggIntegrateFn = std::function<std::vector<Reflection>(
    const std::vector<Reflection> &predicted, size_t npredicted, int64_t image_number)>;

class IndexAndRefine {
    // When false, the current image's result is still returned via the outgoing message, but the
    // whole-run integration_outcome vector is not retained (viewer live/interactive use, which never
    // scales the accumulated run). rugnux/receiver keep it true so ScaleAllImages/merge have the data.
    const bool retain_outcomes_;
    const bool real_time;   // see the constructor
    const DiffractionExperiment& experiment;

    std::optional<CrystalLattice> indexed_lattice;

    std::optional<GoniometerAxis> axis_;

    IndexerThreadPool *indexer_;
    std::unique_ptr<RotationIndexer> rotation_indexer;
    RotationIndexerCounter rotation_indexer_counter;

    struct IndexingOutcome {
        std::optional<CrystalLattice> lattice_candidate;
        std::vector<CrystalLattice> extra_lattice_candidates;
        std::vector<Coord> extra_lattice_rotations;
        DiffractionExperiment experiment;
        LatticeMessage symmetry{
            .centering = 'P',
            .niggli_class = 0,
            .crystal_system = gemmi::CrystalSystem::Triclinic
        };
        bool beam_center_updated = false;

        explicit IndexingOutcome(const DiffractionExperiment& experiment_ref)
            : experiment(experiment_ref) {}
    };

    mutable std::mutex reflections_mutex;
    // Holds the retained outcomes' reflections; declared before them, so it outlives them.
    ReflectionArena reflection_arena;
    std::vector<IntegrationOutcome> integration_outcome;
    std::vector<float> mosaicity;
    // Optional per-frame mosaicity used for Bragg prediction, indexed by image number. When set (the
    // second pass of the rotation two-pass), it overrides the per-image spot-shape estimate so prediction
    // uses the frame-order-SMOOTHED mosaicity that RotationScaleMerge already fitted in the first pass,
    // rather than re-deriving it from scratch.
    std::vector<float> prediction_mosaicity_override_;
    // Predict every node of the lattice, ignoring the centring absences of a fixed space group. Set
    // for the rotation two-pass GEOMETRY pre-pass, whose job is to measure the detector geometry from
    // spot positions and whose intensities are thrown away: rejecting the absences there costs it half
    // its events and buys nothing. Measured with an I-centred group fixed - the pre-pass fitted a
    // different error model (ISa 7.8 -> 3.6), post-refined the distance 119 um away, and the second
    // pass re-indexed 49 of 60 frames instead of 60.
    bool predict_all_centring_nodes_ = false;
    // The lattice centring the last prediction actually ran in: 'P' unless a user-fixed space group
    // made it reject that group's centring absences. What a caller needs to know whether the
    // integration covers every node of the lattice - the P1 cross-check merge does.
    std::atomic<char> prediction_centring_ = 'P';
    // Whether the outgoing message carries its own copy of the integrated reflections. The per-image
    // file writer and the online stream are the only readers of it - the whole-run scaling/merge reads
    // the retained outcome instead - and it is a copy of every reflection of every image, so a caller
    // that writes no per-image file switches it off.
    bool keep_reflections_in_message_ = true;
    // Integrate every rotation frame from the sweep's lattice, whatever fraction of its spots lie on it.
    bool integrate_every_frame_ = false;
    std::vector<float> scale_cc;
    std::vector<std::optional<UnitCell> > unit_cells;

    IndexingOutcome DetermineLatticeAndSymmetryRotation(DataMessage &msg);
    IndexingOutcome DetermineLatticeAndSymmetry(DataMessage &msg);
    // Shared indexing path: determine the lattice/symmetry, refine geometry, and run AnalyzeIndexing.
    // Returns the outcome (ready for integration) when the frame indexes, nullopt otherwise. Both the
    // real per-image ProcessImage and the first-pass scheme validation go through this, so they cannot
    // diverge.
    std::optional<IndexingOutcome> DetermineRefineAnalyze(DataMessage &msg,
                                                          const SpotFindingSettings &spot_finding_settings,
                                                          int64_t *spots_on_lattice = nullptr);
    void RefineGeometryIfNeeded(DataMessage &msg, IndexingOutcome &outcome);
    void QuickPredictAndIntegrate(DataMessage &msg,
                                 const SpotFindingSettings &spot_finding_settings,
                                 BraggPrediction &prediction,
                                 const BraggIntegrateFn &integrate,
                                 const IndexingOutcome &outcome);

    std::unique_ptr<ReindexAmbiguityResolver> reindex_resolver;
    void ScaleImage(DataMessage &msg, IntegrationOutcome& outcome);

    std::optional<float> RotationAngle(int64_t image) const;   // mid-exposure angle for the indexer
public:
    // real_time: bound the geometry refinements - the per-image one here and the candidate-cell ones
    // in the rotation indexer - by WALL CLOCK, as online acquisition must, it having a real budget.
    // Offline (rugnux, the viewer) passes false and they are bounded by iteration count instead, so the
    // same file reprocesses to the same answer regardless of what else the machine was doing.
    IndexAndRefine(const DiffractionExperiment &x, IndexerThreadPool *indexer, bool retain_outcomes = true,
                   bool real_time = false);

    void AddImageToRotationIndexer(DataMessage &msg);
    void ForceRotationIndexerLattice(const CrystalLattice& lattice);
    void ForceRotationIndexerResult(const RotationIndexerResult& result);
    // Supply a per-frame (by image number) mosaicity for prediction, overriding the per-image estimate.
    void SetPredictionMosaicityOverride(std::vector<float> mosaicity_per_frame) {
        prediction_mosaicity_override_ = std::move(mosaicity_per_frame);
    }
    // Predict the centring-absent reflections too, even with a fixed space group - see the member.
    void PredictAllCentringNodes(bool on) { predict_all_centring_nodes_ = on; }
    // The centring prediction ran in - see the member.
    char GetPredictionCentring() const { return prediction_centring_; }
    // Whether the outgoing message keeps its own copy of the reflections - see the member.
    void KeepReflectionsInMessage(bool on) { keep_reflections_in_message_ = on; }
    // See the member.
    void IntegrateEveryFrame(bool on) { integrate_every_frame_ = on; }

    // Returns whether the frame indexed (a lattice was found and refined). Integration, when it runs,
    // is a further step gated on quick_integration.
    void ProcessImage(DataMessage &msg, const SpotFindingSettings &settings,
                      BraggPrediction &prediction, const BraggIntegrateFn &integrate);
    // Index a single frame (no integration) with the current forced rotation lattice; used to score
    // first-pass sampling schemes on the real per-image path. Returns whether the frame indexed.
    bool IndexFrameOnly(DataMessage &msg, const SpotFindingSettings &settings);
    // How many of this frame's spots lie on the forced rotation lattice, whatever the per-frame gate
    // then makes of them. The first pass's acceptance test counts spots over the validation set rather
    // than frames, so it needs the count from the sparse frames IndexFrameOnly refuses too.
    int64_t CountSpotsOnLattice(DataMessage &msg, const SpotFindingSettings &settings);
    IndexAndRefine& ReferenceIntensities(std::vector<MergedReflection> &reference);

    ScalingResult ScaleAllImages(const std::vector<MergedReflection> &reference, size_t nthreads = 0);

    std::optional<RotationIndexerResult> FinalizeRotationIndexing();

    // Re-refine a lattice with its metric constrained to one crystal system, against the rotation
    // indexing's own accumulated spots - for a caller that has adopted a symmetry the indexing never
    // refined under. Empty for a run with no rotation indexer. See RotationIndexer::RefineConstrained.
    [[nodiscard]] std::optional<RotationIndexer::ConstrainedRefit>
    RefineRotationLatticeConstrained(const CrystalLattice &latt, gemmi::CrystalSystem system) const;

    std::optional<UnitCell> GetConsensusUnitCell() const;

    // Not thread safe, need to be run after processing is all done
    const std::vector<float> &GetImageCC() const;
    const std::vector<std::optional<UnitCell> > &GetUnitCells() const;

    std::vector<IntegrationOutcome> &GetIntegrationOutcome();
    const std::vector<IntegrationOutcome> &GetIntegrationOutcome() const;
};