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 <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;
    const DiffractionGeometry geom_;

    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;
    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_;
    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);
    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 per-image geometry refinement by WALL CLOCK, as online acquisition must -
    // it has a real per-image budget. Offline (rugnux, the viewer) passes false and the refinement is
    // 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);
    }

    // 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);
    IndexAndRefine& ReferenceIntensities(std::vector<MergedReflection> &reference);

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

    std::optional<RotationIndexerResult> FinalizeRotationIndexing();

    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;
};