Jungfraujoch/image_analysis/rotation_indexer/RotationIndexer.h

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

#pragma once

#include <memory>
#include <vector>
#include <mutex>

#include "../../common/DiffractionSpot.h"
#include "../../common/DiffractionExperiment.h"
#include "../indexing/IndexerThreadPool.h"
#include "../lattice_search/LatticeSearch.h"

struct RotationIndexerResult {
    CrystalLattice lattice;
    std::vector<CrystalLattice> extra_lattices;
    LatticeSearchResult search_result;
    DiffractionGeometry geom;
    std::optional<GoniometerAxis> axis;
    // The same lattice refined WITHOUT the metric symmetry (triclinic P), when there is one to
    // compare against. The Bravais class is decided on the UNREFINED FFT candidate against a fixed
    // angular tolerance, so a pseudo-symmetric cell can be promoted a class too far and the
    // constraint then snaps a real angle to the ideal one; here the caller can settle that on a
    // statistic the accumulated-spot fraction is too blunt for. Null on the alternative itself.
    std::shared_ptr<RotationIndexerResult> unconstrained;
};

class RotationIndexer {
    mutable std::mutex m;
    const DiffractionExperiment& experiment;

    constexpr static size_t max_spots = 10000;
    constexpr static size_t DEFAULT_MAX_SPOTS_PER_IMAGE = 200;
    // How many of each image's spots join the accumulated cloud, strongest first and ring spots last
    // (FilterSpotsByCount has already ordered them). The default is the bound this always had; the
    // caller lowers it to seed the search on the strongest few, exactly as the per-frame indexer
    // escalates 30 -> 80 -> all, for a pattern where the deep list is mostly not this crystal.
    size_t max_spots_per_image = DEFAULT_MAX_SPOTS_PER_IMAGE;

    const bool index_ice_rings;
    const bool real_time;   // see the constructor

    std::vector<std::vector<SpotToSave>> v_;
    // Per-image rotation angle (mid-exposure, deg) supplied by the caller; falls back to the
    // goniometer evaluated at the image index when not given.
    std::vector<std::optional<float>> angle_deg_;

    std::optional<GoniometerAxis> axis_;
    const DiffractionGeometry geom_;
    DiffractionGeometry updated_geom_;
    LatticeSearchResult search_result_;
    std::vector<CrystalLattice> extra_lattices_;
    std::shared_ptr<RotationIndexerResult> unconstrained_;

    IndexerThreadPool &indexer_;

    size_t accumulated_spots = 0;

    std::optional<CrystalLattice> indexed_lattice;
    std::optional<std::string> indexer_error_;

    // The accumulated per-image spots as one reciprocal-space cloud, back-rotated to phi = 0 under the
    // given geometry and axis. The caller holds the lock.
    [[nodiscard]] std::vector<Coord> AccumulateReciprocal(const DiffractionGeometry &g,
                                                          const std::optional<GoniometerAxis> &ax) const;
public:
    // real_time: bound the candidate-cell refinements by WALL CLOCK, as online acquisition must - it
    // has 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 cell regardless of what else the machine
    // was doing at the time.
    RotationIndexer(const DiffractionExperiment& x, IndexerThreadPool& indexer, bool real_time = false);
    // Keep only this many spots per image in the accumulated cloud (strongest, non-ring first).
    void MaxSpotsPerImage(size_t n) { max_spots_per_image = n > 0 ? n : DEFAULT_MAX_SPOTS_PER_IMAGE; }
    // angle_deg is the image's mid-exposure rotation angle; if omitted, the goniometer angle at
    // `image` is used (only valid when `image` is the goniometer's own image index).
    void ProcessImage(int64_t image, const std::vector<SpotToSave>& spots,
                      std::optional<float> angle_deg = std::nullopt);
    void RunIndexing();
    std::optional<RotationIndexerResult> GetLattice() const;
    // What the indexer threw on the last RunIndexing(), if it threw. Empty when it simply found no
    // lattice: a caller that reports the failure to the user needs the two apart, one being about
    // the crystal and the other about the machine.
    [[nodiscard]] std::optional<std::string> GetIndexerError() const;
    void ForceLattice(const CrystalLattice& lattice);
    // Adopt a complete result (lattice + refined geometry/axis/symmetry) found by another indexer -
    // used when several first-pass sampling schemes are run and the best is selected.
    void ForceResult(const RotationIndexerResult& result);
    // True once the accumulation buffer is full (used to stop feeding a consecutive-frame scheme).
    [[nodiscard]] bool AccumulationFull() const;

    // A lattice re-refined with its metric CONSTRAINED to one crystal system, and what the constraint
    // costs on the spots that decided the free one.
    struct ConstrainedRefit {
        CrystalLattice lattice;
        float indexed_fraction = 0.0f;          // the refined constrained lattice, over the accumulated cloud
        float indexed_fraction_before = 0.0f;   // the lattice handed in, over the same cloud
    };
    // Re-refine `latt` under the metric constraint of `system`, against the same accumulated cloud the
    // indexing ran on and at the geometry and axis it settled on. For a caller that has ADOPTED a
    // symmetry this indexing never refined under: the intensities can confirm a two-fold the spot
    // POSITIONS never offered - the Bravais class is decided on the unrefined candidate, and the free
    // refinement can walk into a class afterwards - and such a caller then holds a cell in a group's
    // setting whose metric is still the free fit's. The detector and the axis are held: by the time
    // that question is asked the reflections have been integrated at that geometry, so what is being
    // asked is which cell the spots support under the constraint, not where the detector is. The BASIS
    // is held too - this never reindexes - so `latt` has to arrive in the setting the constraint is
    // written in (unique axis b for monoclinic, as XtalResidual parameterises it); a caller that cannot
    // guarantee that has to check the answer against its own group.
    [[nodiscard]] std::optional<ConstrainedRefit> RefineConstrained(const CrystalLattice &latt,
                                                                    gemmi::CrystalSystem system) const;
};