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

#include "../../common/JFJochMath.h"
#include "RotationIndexer.h"

#include "../geom_refinement/XtalOptimizer.h"
#include "../indexing/FFTIndexer.h"
#include "../lattice_search/LatticeSearch.h"
#include "../indexing/MultiLatticeSearch.h"

#include <future>

namespace {
    // Sub-cell override thresholds used in candidate selection to undo a spurious axis doubling:
    // a later candidate replaces the chosen cell when it is smaller by more than this volume ratio
    // (a doubling is 2x, well past 1.5) and indexes within this fraction slack of it. The slack is
    // far below the indexed-fraction gap a real superstructure opens between its true cell and its
    // sub-cell, so genuine large cells are kept.
    constexpr float ROT_SUBCELL_VOLUME_RATIO = 1.5f;
    constexpr float ROT_SUBCELL_FRAC_SLACK = 0.02f;

    // How much better a lower-symmetry SETTING of an already-chosen lattice has to index before it is
    // taken (see the selection below). A subgroup setting holds fewer cell parameters fixed, so it can
    // never index less - only a decisively better fit is evidence that the higher symmetry is wrong.
    constexpr float ROT_SUBGROUP_FRAC_RATIO = 1.5f;

    // The same, for a candidate whose primitive cell is a near-integer MULTIPLE of the chosen one.
    // Multiplying an axis halves that reciprocal spacing, so the multiple has a lattice point
    // wherever its sub-cell has one and another in between: it collects spots the sub-cell leaves
    // unindexed, for reasons that have nothing to do with the crystal. The indexed fraction is
    // biased in its favour, so a small lead over the sub-cell is not evidence, and deciding the
    // pair on the ordinary margin leaves it turning on the last bits of that fraction - one
    // crystal here separated the true cell from a spurious 5x supercell by 0.003, little enough
    // that the build's -march flags settled it, and the supercell merged to an R-free of 0.58
    // (i.e. noise). A real superstructure's satellite rows are a large share of its spots and
    // clear this ratio comfortably - but see the limitation noted at the comparison itself.
    constexpr float ROT_SUPERCELL_FRAC_RATIO = 1.5f;
    // How far from a whole number the volume ratio may sit and still count as an axis multiple.
    constexpr double ROT_SUPERCELL_INTEGER_TOL = 0.15;

    // Iteration bound for the candidate-cell refinements when the run is not real-time. Without one,
    // XtalOptimizerData falls back to its wall-clock bound and the cell a candidate refines to - and
    // therefore which candidate wins - depends on how busy the machine was. Same value as Ceres' own
    // default iteration limit, so it only bites where the clock was biting before.
    constexpr int ROT_REFINE_ITERATIONS = 50;

    // Order of the lattice point group, so "lower symmetry" is a well-defined comparison
    // (gemmi's enum orders Trigonal after Tetragonal, which have 6 and 8 rotations).
    int LatticePointGroupOrder(gemmi::CrystalSystem s) {
        switch (s) {
            case gemmi::CrystalSystem::Monoclinic:   return 2;
            case gemmi::CrystalSystem::Orthorhombic: return 4;
            case gemmi::CrystalSystem::Trigonal:     return 6;
            case gemmi::CrystalSystem::Tetragonal:   return 8;
            case gemmi::CrystalSystem::Hexagonal:    return 12;
            case gemmi::CrystalSystem::Cubic:        return 24;
            default:                                 return 1;   // Triclinic
        }
    }

    // Re-express a primitive hexagonal/trigonal lattice in the conventional hexagonal setting
    // (a = b, gamma = 120). The Niggli-reduced primitive cell carries the two equal-length axes
    // at gamma = 60; replacing b with b - a opens that angle to 120 without changing the lattice.
    CrystalLattice HexagonalConventional(CrystalLattice latt) {
        latt.ReorderABEqual();                       // put the equal-length pair in a, b
        Coord a = latt.Vec0(), b = latt.Vec1(), c = latt.Vec2();
        if (angle_deg(a, b) < 90.0f)
            b -= a;
        return CrystalLattice(a, b, c);              // constructor fixes handedness
    }

    bool IsHexagonalSystem(gemmi::CrystalSystem s) {
        return s == gemmi::CrystalSystem::Trigonal || s == gemmi::CrystalSystem::Hexagonal;
    }

    // The hexagonal lattice metric (two equal axes at 60/120 deg, both perpendicular to the third) is
    // also satisfied by its ortho-hexagonal C-centred supercell, so the geometry-keyed LatticeSearch can
    // land there. Detect the hexagonal metric on the reduced PRIMITIVE cell so the de-novo path (no space
    // group to key on) can re-express it in conventional hexagonal axes.
    bool IsMetricallyHexagonal(CrystalLattice latt, float rel_tol = 0.03f, float angle_tol_deg = 3.0f) {
        latt.ReorderABEqual();
        const Coord a = latt.Vec0(), b = latt.Vec1(), c = latt.Vec2();
        const float la = a.Length(), lb = b.Length();
        if (la <= 0.0f || lb <= 0.0f || std::fabs(la - lb) > rel_tol * std::max(la, lb))
            return false;
        const float gab = angle_deg(a, b);
        if (std::fabs(gab - 60.0f) > angle_tol_deg && std::fabs(gab - 120.0f) > angle_tol_deg)
            return false;
        return std::fabs(angle_deg(a, c) - 90.0f) <= angle_tol_deg &&
               std::fabs(angle_deg(b, c) - 90.0f) <= angle_tol_deg;
    }

    // Fraction of the accumulated reciprocal-space spots that a lattice indexes to near-integer
    // Miller indices within tol. Comparing a symmetry-constrained refinement against an
    // unconstrained (triclinic) one is a data-driven test for a false promotion: a wrong
    // higher-symmetry constraint snaps a pseudo cell onto ideal angles and misplaces most spots.
    float IndexedFraction(const CrystalLattice &latt, const std::vector<Coord> &coords, float tol) {
        if (coords.empty())
            return 0.0f;
        const Coord a = latt.Vec0(), b = latt.Vec1(), c = latt.Vec2();
        const float tol_sq = tol * tol;
        size_t indexed = 0;
        for (const Coord &s : coords) {
            const float dh = a * s - std::round(a * s);   // Coord operator* = dot product = Miller index
            const float dk = b * s - std::round(b * s);
            const float dl = c * s - std::round(c * s);
            if (dh * dh + dk * dk + dl * dl < tol_sq)
                ++indexed;
        }
        return static_cast<float>(indexed) / static_cast<float>(coords.size());
    }
}

RotationIndexer::RotationIndexer(const DiffractionExperiment &x, IndexerThreadPool &indexer,
                                 bool real_time)
    : experiment(x),
      index_ice_rings(x.GetIndexingSettings().GetIndexIceRings()),
      real_time(real_time),
      v_(experiment.GetImageNum()),
      angle_deg_(experiment.GetImageNum()),
      axis_(x.GetGoniometer()),
      geom_(x.GetDiffractionGeometry()),
      updated_geom_(geom_),
      indexer_(indexer) {
}

void RotationIndexer::RunIndexing() {
    std::unique_lock ul(m);
    if (!axis_)
        return;

    std::vector<Coord> coords;
    coords.reserve(accumulated_spots);   // the exact count held; the per-image cap may be SIZE_MAX

    for (int i = 0; i < v_.size(); i++) {
        const float angle_deg = angle_deg_[i].value_or(axis_->GetAngle_deg(i) + axis_->GetWedge_deg() / 2.0f);
        const auto rot = axis_->GetTransformationAngle(angle_deg);

        for (const auto &s: v_[i])
            coords.emplace_back(rot * s.ReciprocalCoord(geom_));
    }
    const auto indexer_result = indexer_.Run(experiment, coords);
    indexer_error_ = indexer_result.error;

    if (!indexer_result.lattice.empty() && indexer_result.lattice[0].CalcVolume() > 1.0) {
        DiffractionExperiment experiment_copy(experiment);
        const float index_tol = experiment.GetIndexingSettings().GetTolerance();
        const auto orig_axis = axis_;

        // Map an FFT candidate cell to its Bravais lattice. Re-express a metrically-hexagonal cell in
        // conventional hexagonal axes (LatticeSearch can land on the ortho-hexagonal C setting) so the
        // 3-fold is not hidden from scaling.
        //
        // A user-fixed space group is deliberately NOT stamped on here. The group names the symmetry;
        // it does not say which basis the FFT candidate came back in, and the conventional cell above
        // was reduced for whatever class the METRIC matched. Relabelling that cell with the group's
        // system and centring leaves the constrained refine snapping the wrong angles to the ideal
        // ones: measured, a C-centred orthorhombic cell relabelled primitive monoclinic indexed 1 of
        // 60 validation frames and an F-cubic one relabelled trigonal indexed 0 of 60, where the same
        // frames index at 36/60 and 51/60 without a group. The group is applied where it belongs - to
        // the scaling and the merge.
        auto build_sr = [&](const CrystalLattice &cand) -> LatticeSearchResult {
            auto ls = LatticeSearch(cand);
            if (!IsHexagonalSystem(ls.system) && IsMetricallyHexagonal(ls.primitive_reduced)) {
                ls.conventional = HexagonalConventional(ls.primitive_reduced);
                ls.system = gemmi::CrystalSystem::Hexagonal;
                ls.centering = 'P';
            }
            return ls;
        };

        // Re-accumulate the reciprocal spots under a refined geometry/axis, to score a refined cell.
        auto accumulate = [&](const DiffractionGeometry &g, const std::optional<GoniometerAxis> &ax) {
            return AccumulateReciprocal(g, ax);
        };

        // The FFT offers a few candidate cells (its best reduction plus, for large/elongated cells, a
        // widened alternative). Fully refine each and keep the one that indexes the most spots AFTER
        // geometry refinement - the pre-refinement fraction is not a reliable discriminator (an
        // incorrect larger cell can fit more of the un-refined accumulated spots than the correct one).
        // Twelve, not four. The fraction that orders the candidates is the unreliable one named
        // above, so the correct cell is not always in the first few: on a crystal whose shortlist
        // carries near-degenerate reductions, three of the first four slots went to cells that
        // cannot exist. Refining a candidate is cheap next to the pass that produced it.
        const size_t n_try = std::min<size_t>(indexer_result.lattice.size(), 12);

        // Bound the axis lengths just above the found cell so a free (triclinic) refine cannot drift
        // onto a pseudo-translation / modulation supercell (a modulated crystal whose satellites
        // define a ~4x period would otherwise inflate one axis to the max-length clamp).
        auto make_data = [&](const CrystalLattice &latt, gemmi::CrystalSystem sys, float length_bound_A) {
            XtalOptimizerData d{
                .geom = experiment_copy.GetDiffractionGeometry(),
                .latt = latt,
                .crystal_system = sys,
                .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
                .max_length_A = length_bound_A,
                // Match the indexers' [30,150] deg bound so a monoclinic beta outside [60,120]
                // (e.g. beta>120) is refined, not clamped to the boundary.
                .min_angle_deg = 30.0f,
                .max_angle_deg = 150.0f,
                .refine_beam_center = true,
                .refine_detector_angles = true,
                .refine_rotation_axis = true,
                .index_ice_rings = experiment.GetIndexingSettings().GetIndexIceRings(),
                .max_iterations = real_time ? 0 : ROT_REFINE_ITERATIONS,
                .axis = orig_axis
            };
            if (d.crystal_system == gemmi::CrystalSystem::Trigonal)
                d.crystal_system = gemmi::CrystalSystem::Hexagonal;
            if (d.crystal_system == gemmi::CrystalSystem::Monoclinic)
                d.latt.ReorderMonoclinic();
            return d;
        };

        // Refine the FFT candidates. Each candidate is independent, and within a candidate the
        // metric-symmetry solve and the de-novo triclinic pseudo-symmetry solve are independent too,
        // so refine all of them (up to ~8 solves) at once - these Ceres refinements are the dominant
        // first-pass cost. Each solve runs Ceres on a few cores. Selection stays serial and in
        // candidate order below, so the outcome is identical to refining them one by one.
        constexpr int kCeresThreads = 4;

        // Seed each candidate serially (cheap: LatticeSearch + setup), then solve them in parallel.
        struct CandidateWork {
            bool viable = false;
            LatticeSearchResult sr;
            XtalOptimizerData constrained;
            bool has_tri = false;
            XtalOptimizerData tri;
        };
        std::vector<CandidateWork> work(n_try);
        for (size_t ci = 0; ci < n_try; ci++) {
            const CrystalLattice &cand = indexer_result.lattice[ci];
            if (cand.CalcVolume() <= 1.0)
                continue;
            CandidateWork &w = work[ci];
            w.sr = build_sr(cand);
            const auto conv_uc = w.sr.conventional.GetUnitCell();
            const float length_bound_A = 1.2f * static_cast<float>(std::max({conv_uc.a, conv_uc.b, conv_uc.c}));
            w.constrained = make_data(w.sr.conventional, w.sr.system, length_bound_A);
            // Pseudo-symmetry guard: also refine unconstrained (triclinic) on the primitive cell. Run it
            // with a user-fixed space group too. The metric promotion the constrained refine acts on is
            // decided by the geometry, not by the group, so it can be false whether or not a group was
            // given - and without this cell there is nothing to catch it with, which is how a fixed
            // group turned crystals the de-novo path indexes at 60/60 into runs that index none.
            w.has_tri = (w.sr.system != gemmi::CrystalSystem::Triclinic);
            if (w.has_tri)
                w.tri = make_data(w.sr.primitive_reduced, gemmi::CrystalSystem::Triclinic, length_bound_A);
            w.viable = true;
        }

        // Refine (constrained metric solve + score by the refined-geometry indexed fraction, the
        // reliable discriminator). Runs on its own thread per solve.
        struct Solved { bool ok = false; float frac = 0.0f; XtalOptimizerData data; };
        auto solve = [&](XtalOptimizerData d) -> Solved {
            const bool ok = XtalOptimizer(d, v_, kCeresThreads);
            const float frac = ok ? IndexedFraction(d.latt, accumulate(d.geom, d.axis), index_tol) : 0.0f;
            return {ok, frac, std::move(d)};
        };
        std::vector<std::future<Solved>> constrained_f(n_try), tri_f(n_try);
        for (size_t ci = 0; ci < n_try; ci++) {
            if (!work[ci].viable)
                continue;
            constrained_f[ci] = std::async(std::launch::async, [&, ci] { return solve(work[ci].constrained); });
            if (work[ci].has_tri)
                tri_f[ci] = std::async(std::launch::async, [&, ci] { return solve(work[ci].tri); });
        }

        // Assemble and select serially, in candidate order - identical to refining them one by one.
        float best_frac = -1.0f;
        float best_prim_vol = 0.0f;
        bool have_best = false;
        size_t best_ci = 0;
        XtalOptimizerData best_data;
        LatticeSearchResult best_sr;
        std::shared_ptr<RotationIndexerResult> best_alt;
        for (size_t ci = 0; ci < n_try; ci++) {
            if (!work[ci].viable)
                continue;

            Solved c = constrained_f[ci].get();
            bool ok = c.ok;
            float frac = c.frac;
            XtalOptimizerData data = std::move(c.data);
            LatticeSearchResult sr = work[ci].sr;

            // Adopt the free triclinic cell only if it indexes CLEARLY more than the constrained cell -
            // a false promotion (a near-90 pseudo cell forced to ideal angles + a bogus centering)
            // misplaces most reflections (measured indexed-fraction ratio ~0.1), whereas genuine higher
            // symmetry (incl. R-centred) indexes comparably (ratio ~0.7). Preferring the constrained
            // cell on a near-tie keeps the real symmetry/centering; the intensities settle the final
            // space group.
            // When the guard leaves the constrained cell in place, hand the refined triclinic cell to
            // the caller instead of dropping it: the accumulated-spot fraction separates a false
            // promotion from genuine symmetry by less than 2x on a lattice that is pseudo-symmetric
            // to a few tenths of a degree, while the caller's per-frame validation separates the same
            // pair by more than 20x.
            std::shared_ptr<RotationIndexerResult> tri_alt;
            if (work[ci].has_tri) {
                Solved t = tri_f[ci].get();
                auto as_triclinic = [](LatticeSearchResult s) {
                    s.system = gemmi::CrystalSystem::Triclinic;
                    s.centering = 'P';
                    s.conventional = s.primitive_reduced;
                    s.reindex = gemmi::Mat33(1, 0, 0, 0, 1, 0, 0, 0, 1);
                    return s;
                };
                if (t.ok && t.frac > 0.3f && frac < 0.5f * t.frac) {
                    sr = as_triclinic(sr);
                    data = std::move(t.data);
                    ok = true;
                    frac = t.frac;
                } else if (t.ok) {
                    tri_alt = std::make_shared<RotationIndexerResult>(RotationIndexerResult{
                        .lattice = t.data.latt,
                        .search_result = as_triclinic(sr),
                        .geom = t.data.geom,
                        .axis = t.data.axis,
                    });
                }
            }

            if (!ok)
                continue;

            // Prefer the indexer's earlier (primary) candidate; adopt a later one only if it indexes
            // clearly more AND indexes reasonably well in absolute terms. The absolute floor stops a
            // marginally-higher alternative from displacing the primary when both index poorly (e.g. a
            // twin, where the accumulated-spot fraction is a noisy proxy) - only a decisively better
            // cell (a superstructure's true cell vs its sublattice) takes over.
            // Displace the current best when the candidate indexes clearly more, OR when it is a
            // genuine sub-cell: a meaningfully smaller cell that still indexes at least as many spots.
            // The sub-cell branch unmasks a spurious supercell (axis doubling): the primitive cell
            // always indexes >= its integer multiple, so a doubled cell that wins ci-order by the
            // hysteresis margin is overridden by its own primitive. A real superstructure's true
            // (larger) cell indexes MORE than its sub-cell and is kept by the clearly-more branch;
            // twin lattices share the cell volume, so this never disturbs twin selection.
            // Volumes are compared PRIMITIVE. A centred conventional cell is an exact integer multiple
            // of its primitive one, so two settings of the same lattice differ by that factor and
            // conventional volumes read a mere change of setting as a sub-cell. And two such settings
            // are not comparable on the indexed fraction either: the lower-symmetry one holds fewer
            // cell parameters fixed, so it can only index more. An F-cubic lattice contains an
            // I-tetragonal cell of the same volume, and refining that cell frees the c/a ratio the
            // cubic one holds at sqrt(2), buying back the spots a fraction of a percent of strain had
            // put out of tolerance. A slightly higher fraction is therefore no evidence against the
            // higher symmetry - only a decisively better fit is.
            const float cand_prim_vol = std::abs(data.latt.ToPrimitive(sr.centering).CalcVolume());
            const bool lower_symmetry_setting = have_best
                    && LatticePointGroupOrder(sr.system) < LatticePointGroupOrder(best_sr.system)
                    && std::abs(cand_prim_vol - best_prim_vol) < 0.05f * best_prim_vol;
            // A near-integer volume multiple of the incumbent - see ROT_SUPERCELL_FRAC_RATIO.
            const double vol_ratio = (have_best && best_prim_vol > 1.0f && cand_prim_vol > 1.0f)
                                     ? cand_prim_vol / best_prim_vol : 1.0;
            const double vol_nearest = std::round(vol_ratio);
            const bool integer_supercell = vol_nearest >= 2.0
                    && std::abs(vol_ratio - vol_nearest) < ROT_SUPERCELL_INTEGER_TOL;
            // NOTE, and it is a real limitation: `frac > RATIO * best_frac` cannot be satisfied at
            // all once best_frac exceeds 1/RATIO - above 0.667 for a ratio of 1.5, which is ordinary
            // for good rotation data. So on data that indexes well these two guards do not merely
            // raise the bar, they close the branch: no axis multiple and no lower-symmetry setting
            // can displace the incumbent however much better it fits. A genuine superstructure whose
            // satellite rows the sub-cell misses is therefore kept as its sub-cell, silently.
            //
            // Restating the bar on the fraction left UNINDEXED - the candidate must account for
            // 1/RATIO of what the incumbent missed - is well defined over the whole range and looks
            // like the obvious repair. It was tried and it REGRESSED the 37-crystal battery from
            // 34/37 to 32/37 correct space groups: one C2 lattice fell to P1, and a P2 case went to
            // C222 keeping 2923 of 22440 reflections. The indexed fraction is too noisy a statistic
            // to carry a looser test, so the unreachable-but-safe form stays until the selection is
            // decided on something better than it.
            const bool clearly_more = frac > best_frac + 0.05f && frac > 0.15f
                    && (!lower_symmetry_setting || frac > ROT_SUBGROUP_FRAC_RATIO * best_frac)
                    && (!integer_supercell || frac > ROT_SUPERCELL_FRAC_RATIO * best_frac);
            const bool smaller_subcell = have_best && frac > 0.15f
                                         && frac >= best_frac - ROT_SUBCELL_FRAC_SLACK
                                         && cand_prim_vol < best_prim_vol / ROT_SUBCELL_VOLUME_RATIO;
            if (!have_best || clearly_more || smaller_subcell) {
                best_frac = frac;
                best_prim_vol = cand_prim_vol;
                have_best = true;
                best_data = std::move(data);
                best_sr = sr;
                best_ci = ci;
                best_alt = std::move(tri_alt);
            }
        }

        // Drive the selected candidate to the fit's FIXED POINT. Refinement here is a chain - solve,
        // re-accumulate the reciprocal-space cloud under the refined geometry, solve again - and the
        // selection above calls the solver ONCE per candidate, so what it returns is a point on the way
        // to that fixed point rather than the fixed point itself. One round is enough to rank the
        // candidates; it is not enough to have measured the geometry.
        //
        // What that costs, on a sweep whose 2theta reaches far enough for the detector tilt to be
        // determined at all (long wavelength, short distance): the spindle-parallel tilt comes out at
        // 16 % of an independently measured value after one round and 87 % of it after twenty, and the
        // three sweeps of one crystal at three wavelengths then agree with that external value to
        // 0.02 deg. The data follow - R_meas 0.137 -> 0.119, ISa 7.8 -> 9.2, and the anomalous peak
        // height at the sulphur positions of a known model 7.0 -> 7.8 sigma.
        //
        // It also removes the fit's dependence on where the beam centre started, which is the thing
        // the gauge prior below exists to guard: displacing the starting centre over 8 px moves the
        // refined tilt by 0.0004 deg/px here against 0.0077 deg/px for a single round, on a geometric
        // one-for-one of 0.0710. So the prior stays, and iterating turns it from a prior that pins the
        // answer near the header into a per-step limit that lets the pair walk to their joint optimum -
        // pinning its anchor across the rounds instead leaves the tilt at 0.014 deg and the data worse
        // than doing nothing. The winner only: the other candidates are discarded, and these solves are
        // the dominant first-pass cost. Not in real time, where the budget is wall-clock.
        //
        // The chain is a trajectory, and its last point is not always its best one. Every solve ends
        // by fitting only the spots inside its tightest gate (0.1), so a cell with a direction the
        // data barely speak about - which is what a free cell whose metric is near a Bravais class
        // has - can slide along that direction, pulling a core of spots tighter while the periphery
        // falls out of the fit altogether. Measured on such a crystal: over seventeen rounds the
        // spots inside the tight gate rise from 0.295 to 0.327 while the spots inside the wide one
        // peak at round three (0.754) and fall to 0.730, and it is round three that merges - ISa 11.0
        // against 6.3, R_meas 0.148 against 0.213 - and that agrees with the archived reference cell.
        // Which round a build stopped on used to decide that, since the tilt step the loop tests is
        // scatter at the size of its own bound.
        //
        // So score every round on the WIDE gate the last pass does not fit - the same 0.3 the solve's
        // first pass selects on - and commit the best-scoring round rather than the one the loop
        // happens to stop on. The comparison is between rounds of ONE chain, whose cell moves by a
        // fraction of a percent, so it is not a test a bigger cell can win by being bigger.
        //
        // The score is a COUNT of spots, and a lead of fewer than sqrt(count) of them is smaller than
        // the count's own noise - taking the round on one would hand the answer back to the last
        // bits, which is the defect being removed. So the chain's last round stands unless another
        // beats it by more than that, and the ordinary chain - which settles, and whose rounds then
        // score within a spot or two of each other - commits exactly what it committed before. The
        // bound is generous: both counts are made over the same spot list, so the difference of two
        // rounds carries far less noise than sqrt of either.
        //
        // And the round taken has to be the less distorted lattice as well as the better-fitting one.
        // Ask LatticeSearch for the class of each round's cell - to MEASURE the drift, not to impose
        // the class, which is measured fatal - and take the earlier round only when it matched the
        // SAME class and its metric sits closer to it. Same class, because the deviation is a
        // fraction of the tolerance of whatever class the round matched, so two classes' deviations
        // are not the same quantity - and a round that matched no class reports 0, which is not
        // "undistorted" but "nothing was asserted". Comparing that against a class's deviation reads
        // the absence of a constraint as the absence of distortion, and the veto silently switches
        // off. Measured: it would do that on 11 of 1132 traced chains, four of them against a
        // round that matched no class at all.
        //
        // A lattice's metric symmetry is exact, so among cells that all fit the data the least
        // distorted one is the one the crystal has; a chain that is merely converging does not move
        // its distortion (a symmetry-constrained solve holds it at zero throughout), so this fires on
        // exactly the chains that are sliding out of a class and on no others. It is a veto and not
        // the criterion: the spots have to prefer the round first, so this cannot pull an answer
        // towards a symmetry the data do not support.
        constexpr int ROT_REFINE_OUTER_ROUNDS = 20;
        constexpr double ROT_REFINE_TILT_SETTLED_RAD = 1.0e-5;   // ~0.6 mdeg
        if (have_best && !real_time) {
            auto wide_count = [&](const XtalOptimizerData &d) {
                const auto c = accumulate(d.geom, d.axis);
                return IndexedFraction(d.latt, c, XTAL_OPTIMIZER_WIDE_TOLERANCE) * static_cast<float>(c.size());
            };
            // True when a matched the same Bravais class as b and sits closer to its ideal metric.
            auto less_distorted = [](const LatticeSearchResult &a, const LatticeSearchResult &b) {
                return a.system == b.system && a.centering == b.centering
                       && MetricDeviation(a) < MetricDeviation(b);
            };
            XtalOptimizerData kept = best_data;
            float kept_count = -1.0f;
            float last_count = 0.0f;
            LatticeSearchResult kept_class, last_class;
            for (int r = 0; r < ROT_REFINE_OUTER_ROUNDS; ++r) {
                XtalOptimizerData d = best_data;
                if (!XtalOptimizer(d, v_, kCeresThreads))
                    break;
                const double d1 = std::abs(d.geom.GetPoniRot1_rad() - best_data.geom.GetPoniRot1_rad());
                const double d2 = std::abs(d.geom.GetPoniRot2_rad() - best_data.geom.GetPoniRot2_rad());
                best_data = std::move(d);
                last_count = wide_count(best_data);
                last_class = LatticeSearch(best_data.latt);
                if (last_count >= kept_count) {
                    kept_count = last_count;
                    kept_class = last_class;
                    kept = best_data;
                }
                if (std::max(d1, d2) < ROT_REFINE_TILT_SETTLED_RAD)
                    break;
            }
            if (kept_count > last_count + std::sqrt(last_count) && less_distorted(kept_class, last_class))
                best_data = std::move(kept);
            best_frac = IndexedFraction(best_data.latt, accumulate(best_data.geom, best_data.axis), index_tol);
        }

        if (have_best) {
            search_result_ = best_sr;
            indexed_lattice = best_data.latt;
            updated_geom_ = best_data.geom;
            axis_ = best_data.axis;
            unconstrained_ = std::move(best_alt);
        }

        // Extra (twin) lattices: MultiLatticeSearch derives each rotation by relating the FFT's primary
        // lattice[0] to its near-copies, so only apply it when the chosen cell IS that primary. If a
        // widened alternative won (a superstructure/large cell), lattice[0] is a different (sublattice)
        // metric and its rotations would misorient the chosen cell.
        if (have_best && best_ci == 0 && indexer_result.lattice.size() > 1) {
            auto ml_latt = MultiLatticeSearch(indexer_result.lattice);
            for (auto &l : ml_latt) {
                if (extra_lattices_.size() >= experiment.GetIndexingSettings().GetMaxExtraLattices())
                    break;

                // Ignore lattices oriented by less than 3.0 degree
                if (l.rotation_vector.Length() < 3.0 * PI / 180.0)
                    continue;

                RotMatrix rot(l.rotation_vector.Length(), l.rotation_vector.Normalize());

                XtalOptimizerData data_multi{
                    .geom = experiment_copy.GetDiffractionGeometry(),
                    .latt = indexed_lattice->Multiply(rot),
                    .crystal_system = search_result_.system,
                    .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
                    .refine_beam_center = false,
                    .refine_detector_angles = false,
                    .refine_unit_cell = false,
                    .refine_rotation_axis = false,
                    .index_ice_rings = experiment.GetIndexingSettings().GetIndexIceRings(),
                    .max_iterations = real_time ? 0 : ROT_REFINE_ITERATIONS,
                    .axis = axis_
                };

                // Quick refinement: orientation only. Cell size/angles, beam center,
                // detector angles and rotation axis are all kept from the first lattice.
                // XtalOptimizer always refines orientation; everything else is frozen above.
                XtalOptimizer(data_multi, v_);

                extra_lattices_.push_back(data_multi.latt);
            }
        }
    }
}

void RotationIndexer::ProcessImage(int64_t image, const std::vector<SpotToSave> &spots,
                                   std::optional<float> angle_deg) {
    std::unique_lock ul(m);

    // For non-rotation just ignore the whole procedure
    if (!axis_)
        return;

    // Guard: `image` is a slot in [0, image count); a bad index (e.g. a global number for a subset
    // run) must not corrupt memory.
    if (image < 0 || image >= static_cast<int64_t>(v_.size()))
        return;

    if (accumulated_spots >= max_spots)
        return;

    if (indexed_lattice)
        return;

    angle_deg_[image] = angle_deg;
    v_[image].reserve(spots.size());

    for (const auto &s: spots) {
        if (index_ice_rings || !s.ice_ring)
            v_[image].emplace_back(s);
    }

    // truncate spots, so we don't get above max_spots (total) and max_spots_per_image (for this image)
    size_t max_spots_limit = std::min(max_spots_per_image, max_spots - accumulated_spots);

    if (v_[image].size() > max_spots_limit) {
        std::ranges::nth_element(v_[image], v_[image].begin() + max_spots_limit,
                                 [](const SpotToSave &a, const SpotToSave &b) {
                                     return a.intensity > b.intensity;
                                 }
        );

        v_[image].resize(max_spots_limit);
    }

    accumulated_spots += v_[image].size();
}

std::vector<Coord> RotationIndexer::AccumulateReciprocal(const DiffractionGeometry &g,
                                                         const std::optional<GoniometerAxis> &ax) const {
    std::vector<Coord> c;
    c.reserve(accumulated_spots);
    for (int i = 0; i < v_.size(); i++) {
        const float a = angle_deg_[i].value_or(ax->GetAngle_deg(i) + ax->GetWedge_deg() / 2.0f);
        const auto rot = ax->GetTransformationAngle(a);
        for (const auto &s : v_[i])
            c.emplace_back(rot * s.ReciprocalCoord(g));
    }
    return c;
}

std::optional<RotationIndexer::ConstrainedRefit>
RotationIndexer::RefineConstrained(const CrystalLattice &latt, gemmi::CrystalSystem system) const {
    std::unique_lock ul(m);

    if (!axis_ || accumulated_spots == 0 || latt.CalcVolume() <= 1.0)
        return {};

    const UnitCell uc = latt.GetUnitCell();
    XtalOptimizerData d{
        .geom = updated_geom_,
        .latt = latt,
        .crystal_system = system == gemmi::CrystalSystem::Trigonal ? gemmi::CrystalSystem::Hexagonal
                                                                   : system,
        .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
        .max_length_A = 1.2f * std::max({uc.a, uc.b, uc.c}),
        // The indexers' bound, so a monoclinic beta outside [60,120] is refined and not clamped.
        .min_angle_deg = 30.0f,
        .max_angle_deg = 150.0f,
        // The detector and the axis stay where the indexing left them - see the header.
        .refine_beam_center = false,
        .refine_detector_angles = false,
        .refine_rotation_axis = false,
        .index_ice_rings = index_ice_rings,
        .max_iterations = real_time ? 0 : ROT_REFINE_ITERATIONS,
        .axis = axis_
    };

    // One cloud for both scores: the geometry is held, so the lattice that comes out is scored on the
    // same spots the one that went in is.
    const auto cloud = AccumulateReciprocal(updated_geom_, axis_);
    const float tol = experiment.GetIndexingSettings().GetTolerance();
    const float before = IndexedFraction(latt, cloud, tol);

    if (!XtalOptimizer(d, v_, 4))
        return {};

    return ConstrainedRefit{
        .lattice = d.latt,
        .indexed_fraction = IndexedFraction(d.latt, cloud, tol),
        .indexed_fraction_before = before,
    };
}

std::optional<RotationIndexerResult> RotationIndexer::GetLattice() const {
    std::unique_lock ul(m);

    if (!indexed_lattice)
        return {};
    return RotationIndexerResult{
        .lattice = indexed_lattice.value(),
        .extra_lattices = extra_lattices_,
        .search_result = search_result_,
        .geom = updated_geom_,
        .axis = axis_,
        .unconstrained = unconstrained_,
    };
}

std::optional<std::string> RotationIndexer::GetIndexerError() const {
    std::unique_lock ul(m);
    return indexer_error_;
}

void RotationIndexer::ForceResult(const RotationIndexerResult &result) {
    std::unique_lock ul(m);
    indexed_lattice = result.lattice;
    extra_lattices_ = result.extra_lattices;
    search_result_ = result.search_result;
    updated_geom_ = result.geom;
    axis_ = result.axis;
    unconstrained_ = result.unconstrained;
}

bool RotationIndexer::AccumulationFull() const {
    std::unique_lock ul(m);
    return accumulated_spots >= max_spots;
}

void RotationIndexer::ForceLattice(const CrystalLattice &lattice) {
    indexed_lattice = lattice;
    const gemmi::SpaceGroup &sg = experiment.GetSpaceGroupOrP1();
    search_result_ = LatticeSearchResult{
        .niggli_class = 0, // Since Niggli class was not searched for, we don't know which one
        .conventional = lattice, // If lattice provided, it is for now primitive == conventional
        .system = sg.crystal_system(),
        .centering = sg.centring_type(),
    };
}