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 <stdexcept>

#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;
    // Set where the free fit walked the detector tilt further from the tilt the indexer was handed
    // than a mounted detector is off square by (ROT_TILT_PRIOR_DEG in RotationIndexer.cpp), and the
    // lattice was therefore refined again with the tilt held where it started and the two compared
    // on the spots they index: the tilt the fit walked to, the two counts, and which won. Refused,
    // geom carries the tilt the indexer started at; confirmed, the walked tilt.
    struct TiltWalk {
        std::array<double, 2> tilt_rad;
        float spots_walked = 0.0f;
        float spots_held = 0.0f;
        bool refused = false;
    };
    std::optional<TiltWalk> tilt_walk;
    // Twinning by reticular merohedry (rugnux's twin-composite step): a second domain whose lattice
    // shares the reflections h of this one with h . t divisible by n, t a vector of this lattice (real
    // space, in the frame of `lattice`, so it holds through any change of basis). Those reflections
    // carry both domains' intensity, and are not integrated.
    struct TwinOverlap {
        Coord t;
        int n = 0;
    };
    std::optional<TwinOverlap> twin_overlap;
};

// What a memo verification throws where a reused result differs from its recomputation - its own
// type, so that no pass that catches a failure and carries on can take it for an ordinary one.
struct MemoMismatch : std::runtime_error {
    using std::runtime_error::runtime_error;
};

class RotationIndexer {
public:
    // RunIndexing is deterministic in what it reads - the accumulated spots and their angles, the
    // geometry and axis, the experiment's indexing settings, cell and space group, and the settings
    // of the pool it indexes with - and a rotation run asks the same question again and again (the
    // canonical pass repeats its rotation-scale probe, a probe repeats its pass's rescue ladder).
    // So what it computes is kept, process-wide, under exactly those inputs, and a second
    // RotationIndexer asking with the same inputs takes the answer. VerifyMemo(true) recomputes anyway
    // and throws on any difference.
    struct IndexingOutcome {
        std::optional<GoniometerAxis> axis;
        DiffractionGeometry updated_geom;
        LatticeSearchResult search_result;
        std::vector<CrystalLattice> extra_lattices;
        std::shared_ptr<RotationIndexerResult> unconstrained;
        std::optional<RotationIndexerResult::TiltWalk> tilt_walk;
        std::optional<CrystalLattice> indexed_lattice;
    };
private:
    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;
    bool verify_memo = false;

    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_;
    std::optional<RotationIndexerResult::TiltWalk> tilt_walk_;
    std::optional<RotationIndexerResult::TwinOverlap> twin_overlap_;   // only ever forced, see ForceResult

    IndexerThreadPool &indexer_;

    size_t accumulated_spots = 0;

    std::optional<CrystalLattice> indexed_lattice;

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

    [[nodiscard]] std::vector<double> InputKey() const;
    [[nodiscard]] IndexingOutcome Outcome() const;
    void RunIndexingLocked();
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; }
    // Recompute on a memo hit and throw if the two differ - the tests' check that the memo is exact.
    void VerifyMemo(bool v) { verify_memo = v; }
    // 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);
    // Throws what the indexer threw, where it failed rather than found no lattice: the two are one
    // about the machine and one about the crystal, and only the second is an answer.
    void RunIndexing();
    std::optional<RotationIndexerResult> GetLattice() 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;
};