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 <array>
#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)>;

// What the supercell probe (IndexAndRefine::ProbeSupercell) reads for one of the eight parity classes
// of the 2a x 2b x 2c primitive supercell in one resolution shell: class 4*(h&1) + 2*(k&1) + (l&1) in
// supercell indices, so class 0 is the adopted lattice and each of the other seven is one index-2
// superstructure. The partiality sums are for the fit I = a + b p over the class: b is the part of
// it that rocks like a Bragg reflection, a what sits at its position whatever the rocking.
struct SupercellProbeClass {
    int64_t n = 0;
    double sum_i = 0.0, sum_i_over_sigma = 0.0;
    double sum_p = 0.0, sum_pp = 0.0, sum_pi = 0.0, sum_ii = 0.0;
    SupercellProbeClass &operator+=(const SupercellProbeClass &o) {
        n += o.n;
        sum_i += o.sum_i; sum_i_over_sigma += o.sum_i_over_sigma;
        sum_p += o.sum_p; sum_pp += o.sum_pp; sum_pi += o.sum_pi; sum_ii += o.sum_ii;
        return *this;
    }
};
// [class][shell], the shells 20-5 A and 5-3 A.
using SupercellProbe = std::array<std::array<SupercellProbeClass, 2>, 8>;

// One parity class of the supercell probe, fitted I = a + b p over partiality p: b (with its standard
// error) is what rocks like a Bragg reflection, a what does not.
struct SupercellProbeFit {
    double mean_i = 0.0, mean_i_over_sigma = 0.0, mean_p = 0.0;
    double a = 0.0, b = 0.0, b_se = 0.0;
};
SupercellProbeFit FitSupercellProbe(const SupercellProbeClass &c);

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;

    // Supercell probe: the image numbers it reads, and what it has read (see ProbeSupercell).
    std::vector<int64_t> supercell_probe_frames_;
    SupercellProbe supercell_probe_{};
    std::optional<CrystalLattice> supercell_probe_primitive_;
    std::mutex supercell_probe_mutex_;
    void ProbeSupercellFrame(const DataMessage &msg, BraggPrediction &prediction,
                             const BraggIntegrateFn &integrate, const IndexingOutcome &outcome,
                             const BraggPredictionSettings &settings);

    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,
                                 const BraggIntegrateFn &probe_integrate);

    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.
    // probe_integrate integrates the supercell probe's reflections (ProbeSupercell); a caller that
    // gives none takes no part in the probe.
    void ProcessImage(DataMessage &msg, const SpotFindingSettings &settings,
                      BraggPrediction &prediction, const BraggIntegrateFn &integrate,
                      const BraggIntegrateFn &probe_integrate = {});

    // Does the lattice each frame is refined on have an index-2 superstructure? On the given frames,
    // after the frame's own integration, predict the 2a x 2b x 2c supercell of its PRIMITIVE lattice
    // (the per-frame refinement and geometry unchanged), integrate it out to 3 A, and sum the eight
    // parity classes in two shells over 20-3 A (SupercellProbe). Class 0 is the lattice itself; the
    // other seven are the seven index-2 superlattices, at most one of which a single doubling fills.
    // The probe's integrations are the caller's to keep out of anything else it counts.
    void ProbeSupercell(std::vector<int64_t> image_numbers);
    [[nodiscard]] SupercellProbe GetSupercellProbe();
    // The primitive lattice the parity classes are indexed in (the last probed frame's; only its metric
    // is meant to be read), so a class can be turned into the cell it would double to.
    [[nodiscard]] std::optional<CrystalLattice> GetSupercellProbePrimitive();
    // 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;
};