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 max_spots_per_image = 200;

    const bool index_ice_rings;

    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;
public:
    RotationIndexer(const DiffractionExperiment& x, IndexerThreadPool& indexer);
    // 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;
    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;
};