Jungfraujoch/image_analysis/IndexAndRefine.cpp

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

#include <limits>
#include <cstdlib>
#include "IndexAndRefine.h"

#include "bragg_integration/CalcISigma.h"
#include "geom_refinement/XtalOptimizer.h"
#include "indexing/AnalyzeIndexing.h"
#include "indexing/FFTIndexer.h"
#include "indexing/MultiLatticeSearch.h"
#include "lattice_search/LatticeSearch.h"
#include "scale_merge/ReindexAmbiguity.h"
#include "scale_merge/ScaleOnTheFly.h"

IndexAndRefine::IndexAndRefine(const DiffractionExperiment &x, IndexerThreadPool *indexer, bool retain_outcomes)
    : index_ice_rings(x.GetIndexingSettings().GetIndexIceRings()),
      retain_outcomes_(retain_outcomes),
      experiment(x),
      geom_(x.GetDiffractionGeometry()),
      indexer_(indexer),
      rotation_indexer_counter(x) {
    if (indexer && x.IsRotationIndexing())
        rotation_indexer = std::make_unique<RotationIndexer>(x, *indexer);
    // Only retain the whole-run per-image reflections when a later scaling/merge pass will read them.
    if (retain_outcomes_)
        integration_outcome.resize(x.GetImageNum());

    mosaicity.resize(x.GetImageNum(), NAN);
    scale_cc.resize(x.GetImageNum(), 0);
    unit_cells.resize(x.GetImageNum());
}

std::optional<float> IndexAndRefine::RotationAngle(int64_t image) const {
    // Mid-exposure rotation angle for image index `image`, matching the angle used for prediction.
    if (const auto g = experiment.GetGoniometer())
        return g->GetAngle_deg(static_cast<float>(image)) + g->GetWedge_deg() / 2.0f;
    return std::nullopt;
}

void IndexAndRefine::AddImageToRotationIndexer(DataMessage &msg) {
    if (rotation_indexer)
        rotation_indexer->ProcessImage(msg.number, msg.spots, RotationAngle(msg.number));
}

IndexAndRefine::IndexingOutcome IndexAndRefine::DetermineLatticeAndSymmetryRotation(DataMessage &msg) {
    IndexingOutcome outcome(experiment);

    if (!rotation_indexer)
        return outcome;

    auto result = rotation_indexer->GetLattice();

    if (!result.has_value()) {
        auto rot_cnt = rotation_indexer_counter.Process(msg.number);
        if (rot_cnt.first)
            rotation_indexer->ProcessImage(msg.number, msg.spots, RotationAngle(msg.number));
        if (rot_cnt.second)
            rotation_indexer->RunIndexing();
        result = rotation_indexer->GetLattice();
    }

    if (result.has_value()) {
        // For rotation indexing, indexing rate is calculated only for frames, where "global" rotation indexing solution was found
        msg.indexing_result = false;

        // get rotated lattice
        auto gon = result->axis;
        if (gon) {
            const float angle_deg = gon->GetAngle_deg(msg.number) + gon->GetWedge_deg() / 2.0f;
            const auto rot_to_image = gon->GetTransformationAngle(-angle_deg);
            outcome.lattice_candidate = result->lattice.Multiply(rot_to_image);

            outcome.extra_lattice_candidates.reserve(result->extra_lattices.size());
            for (const auto &el : result->extra_lattices)
                outcome.extra_lattice_candidates.push_back(el.Multiply(rot_to_image));
        }

        outcome.experiment.BeamX_pxl(result->geom.GetBeamX_pxl())
                .BeamY_pxl(result->geom.GetBeamY_pxl())
                .DetectorDistance_mm(result->geom.GetDetectorDistance_mm())
                .PoniRot1_rad(result->geom.GetPoniRot1_rad())
                .PoniRot2_rad(result->geom.GetPoniRot2_rad())
        .Goniometer(result->axis);
        outcome.symmetry.centering = result->search_result.centering;
        outcome.symmetry.niggli_class = result->search_result.niggli_class;
        outcome.symmetry.crystal_system = result->search_result.system;
    }
    return outcome;
}

IndexAndRefine::IndexingOutcome IndexAndRefine::DetermineLatticeAndSymmetry(DataMessage &msg) {
    auto indexing_start_time = std::chrono::steady_clock::now();

    IndexingOutcome outcome(experiment);

    // Seed the indexer with the strongest few spots first and escalate to more only if that fails.
    // On flooded / noisy frames (XFEL, ice) a lean high-quality seed finds the lattice far more
    // reliably than the full spot list, whose many spurious peaks derail the search; on clean frames
    // the lean seed already works, so nothing is lost. The FULL spot list is still used for geometry
    // refinement and integration downstream, so higher-resolution accuracy is preserved. Cost is
    // ~1 indexer call on frames that index cleanly, up to 3 only on the hard ones. msg.spots is
    // already ordered non-ice-first, strongest-first (FilterSpotsByCount), so the prefix IS the seed.
    const float idx_tol = experiment.GetIndexingSettings().GetTolerance();
    const float idx_tol_sq = idx_tol * idx_tol;
    constexpr float SEED_STOP_FRACTION = 0.9f;   // seed explained this well -> stop escalating
    IndexerResult indexer_result;
    bool any_executed = false;
    float best_frac = -1.0f;
    for (size_t seed_cap : {size_t{30}, size_t{80}, std::numeric_limits<size_t>::max()}) {
        std::vector<Coord> recip;
        recip.reserve(std::min<size_t>(seed_cap, msg.spots.size()));
        for (const auto &i: msg.spots) {
            if (index_ice_rings || !i.ice_ring) {
                recip.push_back(i.ReciprocalCoord(geom_));
                if (recip.size() >= seed_cap)
                    break;
            }
        }
        auto res = indexer_->Run(experiment, recip);
        any_executed |= res.executed;
        if (!res.lattice.empty()) {
            // Keep the seed the lattice explains the largest FRACTION of: a lean clean seed a good
            // lattice indexes almost fully beats a flooded seed it fits only in small part. This
            // auto-selects the lean seed on noisy frames (XFEL) and the full seed where the extra spots
            // are real signal (weak synchrotron) -- no per-dataset setting.
            const Coord a = res.lattice[0].Vec0(), b = res.lattice[0].Vec1(), c = res.lattice[0].Vec2();
            int n = 0;
            for (const auto &q : recip) {
                const float hf = q * a, kf = q * b, lf = q * c;
                const float dh = hf - std::round(hf), dk = kf - std::round(kf), dl = lf - std::round(lf);
                if (dh * dh + dk * dk + dl * dl < idx_tol_sq) ++n;
            }
            const float frac = recip.empty() ? 0.0f : static_cast<float>(n) / recip.size();
            if (frac > best_frac) { best_frac = frac; indexer_result = std::move(res); }
            // A lattice that already explains nearly the whole seed is kept whatever a larger seed
            // returns: the winner is the highest explained FRACTION, and adding weaker spots almost
            // always lowers it. Stop here - this is what keeps clean frames at one indexer call.
            if (frac >= SEED_STOP_FRACTION)
                break;
        }
        if (recip.size() < seed_cap)   // already fed every available spot; a larger cap won't add any
            break;
    }

    if (any_executed)
        msg.indexing_result = false;

    if (!indexer_result.lattice.empty()) {
        auto latt = indexer_result.lattice[0];
        if (latt.CalcVolume() > 1.0) {
            auto sg = experiment.GetGemmiSpaceGroup();
            const auto algorithm = experiment.GetIndexingAlgorithm();
            const bool de_novo = (algorithm == IndexingAlgorithmEnum::FFT
                                  || algorithm == IndexingAlgorithmEnum::FFTW);

            // If space group and cell provided => enforce that symmetry in refinement.
            // If not => detect the symmetry from the lattice.
            if (sg && experiment.GetUnitCell()) {
                outcome.symmetry = LatticeMessage{
                    .centering = sg->centring_type(),
                    .niggli_class = 0,
                    .crystal_system = sg->crystal_system()
                };
                // Place every frame's cell in ONE consistent setting for the whole dataset:
                // mixed axis orders (e.g. [78,78,38] vs [38,78,78]) index the same reflection
                // as different HKLs and cannot be merged. LatticeSearch gives the conventional
                // setting when its detected symmetry agrees with the user's space group. On
                // noisy frames it can instead pick an alternative Bravais setting (e.g. the
                // sqrt2 C-centred description of a primitive tetragonal cell,
                // [78,78,38]->[110,111,38]); there:
                //   - FFBIDX already returns the reference setting (c-last), consistent with the
                //     conventional frames, so its raw lattice is safe -> use it (FFBIDX neutral);
                //   - de-novo indexers (FFT/FFTW) return a Niggli-primitive cell with a DIFFERENT
                //     axis order (c-first) that would corrupt the merge -> reject the frame.
                // niggli_class is left unassigned (0): it needs the primitive cell incl.
                // centering, which LatticeSearch cannot recover from a (possibly centred, e.g.
                // C2) user cell. A proper primitive-cell indexing path (CrystFEL-style) is deferred.
                auto sym_result = LatticeSearch(latt);
                if (sym_result.system == sg->crystal_system())
                    outcome.lattice_candidate = sym_result.conventional;
                else if (!de_novo)
                    outcome.lattice_candidate = latt;
                // else: de-novo + symmetry mismatch -> leave unset, frame is not indexed
            } else {
                auto sym_result = LatticeSearch(latt);
                outcome.symmetry = LatticeMessage{
                    .centering = sym_result.centering,
                    .niggli_class = sym_result.niggli_class,
                    .crystal_system = sym_result.system
                };
                outcome.lattice_candidate = sym_result.conventional;
            }

            // Multi-lattice search for stills: store rotations that map the reference
            // lattice to each accepted extra lattice. Candidates are materialized later
            // in RefineGeometryIfNeeded so they're rooted in the refined (and, for
            // monoclinic, reordered) main lattice.
            if (outcome.lattice_candidate && indexer_result.lattice.size() > 1) {
                auto ml_latt = MultiLatticeSearch(indexer_result.lattice);
                for (auto &ml : ml_latt) {
                    if (outcome.extra_lattice_rotations.size() >= experiment.GetIndexingSettings().GetMaxExtraLattices())
                        break;
                    outcome.extra_lattice_rotations.push_back(ml.rotation_vector);
                    RotMatrix rot(ml.rotation_vector.Length(), ml.rotation_vector.Normalize());
                    outcome.extra_lattice_candidates.push_back(outcome.lattice_candidate->Multiply(rot));
                }
            }

        }
    }

    auto indexing_end_time = std::chrono::steady_clock::now();
    msg.indexing_time_s = std::chrono::duration<float>(indexing_end_time - indexing_start_time).count();

    return outcome;
}

namespace {
// Count spots whose fractional Miller index falls within the indexing tolerance of an integer for a
// given lattice + geometry - the "how well does this model explain the spots" score used by -r multi.
int CountIndexedSpots(const DiffractionGeometry &geom, const CrystalLattice &latt,
                      const std::vector<SpotToSave> &spots, float tol_sq) {
    const Coord a = latt.Vec0(), b = latt.Vec1(), c = latt.Vec2();
    int n = 0;
    for (const auto &s : spots) {
        const Coord recip = s.ReciprocalCoord(geom);
        const float hf = recip * a, kf = recip * b, lf = recip * c;
        const float dh = hf - std::round(hf), dk = kf - std::round(kf), dl = lf - std::round(lf);
        if (dh * dh + dk * dk + dl * dl < tol_sq) ++n;
    }
    return n;
}
}   // namespace

void IndexAndRefine::RefineGeometryIfNeeded(DataMessage &msg, IndexAndRefine::IndexingOutcome &outcome) {
    if (!outcome.lattice_candidate)
        return;

    auto start_time = std::chrono::steady_clock::now();

    XtalOptimizerData data{
        .geom = outcome.experiment.GetDiffractionGeometry(),
        .latt = *outcome.lattice_candidate,
        .crystal_system = outcome.symmetry.crystal_system,
        .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
        // Match the [30,150] deg bound the indexers already use (FFBIDXIndexer, FFT settings):
        // the struct default [60,120] clamps a monoclinic beta outside that window (e.g. a C2
        // beta near 132 deg) to the boundary, corrupting the per-frame cell refinement.
        .min_angle_deg = 30.0f,
        .max_angle_deg = 150.0f,
        .refine_beam_center = true,
        .refine_distance_mm = false,
        .refine_detector_angles = false,
        .refine_unit_cell = !experiment.IsRotationIndexing(),
        // The whole spot list is passed below, not the indexed subset, so on weak images most of what
        // enters the fit at the loose first tolerance is arbitrarily indexed noise. Weight every spot by
        // how strong it is for its resolution so those contribute without dragging the orientation.
        .weight_spots_by_confidence = true,
        .max_time = 0.04 // 40 ms is max allowed time for the operation
    };


    if (outcome.symmetry.crystal_system == gemmi::CrystalSystem::Trigonal)
        data.crystal_system = gemmi::CrystalSystem::Hexagonal;

    switch (experiment.GetIndexingSettings().GetGeomRefinementAlgorithm()) {
        case GeomRefinementAlgorithmEnum::None:
            break;
        case GeomRefinementAlgorithmEnum::OrientationOnly:
            XtalOptimizerRotationOnly(data, msg.spots, 0.2);
            XtalOptimizerRotationOnly(data, msg.spots, 0.1);
            XtalOptimizerRotationOnly(data, msg.spots, 0.05);
            break;
        case GeomRefinementAlgorithmEnum::BeamCenter:
            if (XtalOptimizer(data, {msg.spots})) {
                outcome.experiment.BeamX_pxl(data.geom.GetBeamX_pxl())
                        .BeamY_pxl(data.geom.GetBeamY_pxl());
                outcome.beam_center_updated = true;
            }
            break;
        case GeomRefinementAlgorithmEnum::Flex: {
            // Try all three refinements per image and keep whichever indexes the most spots. Beam+cell
            // refinement helps some stills but diverges on sparse spot lists (few spots, long axes),
            // where it pushes a good lattice out of tolerance; scoring by indexed-spot count lets each
            // image fall back to orientation-only or no refinement when refinement would hurt. Ties
            // prefer less refinement (strict >, not >=) to avoid overfitting.
            const float tol = experiment.GetIndexingSettings().GetTolerance();
            const float tol_sq = tol * tol;

            XtalOptimizerData d_none = data;
            XtalOptimizerData d_orient = data;
            XtalOptimizerRotationOnly(d_orient, msg.spots, 0.2);
            XtalOptimizerRotationOnly(d_orient, msg.spots, 0.1);
            XtalOptimizerRotationOnly(d_orient, msg.spots, 0.05);
            XtalOptimizerData d_beam = data;
            const bool beam_ok = XtalOptimizer(d_beam, {msg.spots});

            const int s_none = CountIndexedSpots(d_none.geom, d_none.latt, msg.spots, tol_sq);
            const int s_orient = CountIndexedSpots(d_orient.geom, d_orient.latt, msg.spots, tol_sq);
            const int s_beam = beam_ok ? CountIndexedSpots(d_beam.geom, d_beam.latt, msg.spots, tol_sq) : -1;

            if (s_beam > s_none && s_beam > s_orient) {
                data = d_beam;
                outcome.experiment.BeamX_pxl(data.geom.GetBeamX_pxl())
                        .BeamY_pxl(data.geom.GetBeamY_pxl());
                outcome.beam_center_updated = true;
            } else if (s_orient > s_none) {
                data = d_orient;
            } else {
                data = d_none;
            }
            break;
        }
    }

    outcome.lattice_candidate = data.latt;

    if (outcome.symmetry.crystal_system == gemmi::CrystalSystem::Monoclinic)
        outcome.lattice_candidate->ReorderMonoclinic();

    // Rebuild extra-lattice candidates from the refined (and possibly reordered) main
    // lattice so they share its cell and obtuse-beta convention.
    if (!outcome.extra_lattice_rotations.empty()) {
        outcome.extra_lattice_candidates.clear();
        outcome.extra_lattice_candidates.reserve(outcome.extra_lattice_rotations.size());
        for (const auto &rv : outcome.extra_lattice_rotations) {
            RotMatrix rot(rv.Length(), rv.Normalize());
            outcome.extra_lattice_candidates.push_back(outcome.lattice_candidate->Multiply(rot));
        }
    }

    // Quick orientation-only refinement of extra lattices (stills path).
    // Cell, beam center, detector geometry are taken from the first lattice.
    if (!experiment.IsRotationIndexing() && !outcome.extra_lattice_candidates.empty()) {
        for (auto &el : outcome.extra_lattice_candidates) {
            XtalOptimizerData data_extra{
                .geom = data.geom,
                .latt = el,
                .crystal_system = data.crystal_system,
                .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
                .refine_beam_center = false,
                .refine_distance_mm = false,
                .refine_detector_angles = false,
                .refine_unit_cell = false,
                .refine_rotation_axis = false,
                .index_ice_rings = experiment.GetIndexingSettings().GetIndexIceRings(),
                .max_time = 0.02
            };
            XtalOptimizerRotationOnly(data_extra, msg.spots, 0.1);
            el = data_extra.latt;
        }
    }

    if (outcome.beam_center_updated) {
        msg.beam_corr_x = data.beam_corr_x;
        msg.beam_corr_y = data.beam_corr_y;
    }

    auto end_time = std::chrono::steady_clock::now();
    msg.refinement_time_s = std::chrono::duration_cast<std::chrono::duration<double>>(end_time - start_time).count();
}

void IndexAndRefine::QuickPredictAndIntegrate(DataMessage &msg,
                                              const SpotFindingSettings &spot_finding_settings,
                                              BraggPrediction &prediction,
                                              const BraggIntegrateFn &integrate,
                                              const IndexAndRefine::IndexingOutcome &outcome) {
    if (!outcome.lattice_candidate)
        return;


    CrystalLattice latt = outcome.lattice_candidate.value();

    // Prediction uses each frame's OWN mosaicity/profile_radius (image-local). We deliberately do NOT
    // smooth them here with a running moving average: it averaged the last N *processed* frames, whose
    // order under the parallel per-image loop is thread-arrival order, making the predicted rocking
    // width - and hence which reflections are integrated - non-deterministic run-to-run. Prediction only
    // decides membership (a reflection on the cutoff contributes ~nothing), so the per-frame value is
    // fine here. The mosaicity smoothing that actually matters - keeping the partialities of one rocking
    // event consistent so they tile the curve and sum toward 1 - is done deterministically in frame
    // order before the 3D combine (RotationScaleMerge), where partiality is recomputed from it.

    float ewald_dist_cutoff = 0.001f;
    if (msg.profile_radius)
        ewald_dist_cutoff = msg.profile_radius.value() * 2.0f;

    if (experiment.GetBraggIntegrationSettings().GetFixedProfileRadius_recipA())
        ewald_dist_cutoff = experiment.GetBraggIntegrationSettings().GetFixedProfileRadius_recipA().value() * 3.0f;

    float wedge_deg = 0.0f;
    float mos_deg = 0.1f;

    if (experiment.GetGoniometer().has_value()) {
        // Full oscillation wedge of one frame; BraggPredictionRot halves it to the +/- half-wedge of the
        // partiality erf pair (Kabsch). Passing the full increment gives a half-wedge of increment/2 -
        // matching ScaleOnTheFly's RotationPartiality, so the predicted partiality is used directly there.
        wedge_deg = experiment.GetGoniometer()->GetWedge_deg();

        if (msg.mosaicity_deg) {
            mos_deg = msg.mosaicity_deg.value();
            mosaicity[msg.number] = mos_deg;
        }
        // Second pass of the rotation two-pass: widen the prediction to the frame-order-smoothed mosaicity
        // that RotationScaleMerge fitted in the first pass. Take the MAX with this frame's own estimate so
        // prediction is never NARROWER than the first pass - a too-narrow smoothed value would otherwise drop
        // reflections and collapse the multiplicity. (A wider value only helps prediction cover the spot.)
        if (msg.number >= 0 && msg.number < static_cast<int64_t>(prediction_mosaicity_override_.size())
            && std::isfinite(prediction_mosaicity_override_[msg.number])
            && prediction_mosaicity_override_[msg.number] > 0.0f) {
            mos_deg = std::max(mos_deg, prediction_mosaicity_override_[msg.number]);
            mosaicity[msg.number] = mos_deg;
        }
    }

    IntegrationOutcome i_outcome{
        .geom = outcome.experiment.GetDiffractionGeometry(),
        .latt = latt,
        .mosaicity_deg = mos_deg,
        .image_scale_cc = msg.image_scale_cc,
    };

    const BraggPredictionSettings settings_prediction{
        .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(),
        .ewald_dist_cutoff = ewald_dist_cutoff,
        .max_hkl = 100,
        // Centering is a hypothesis to confirm, not assume: with no user-fixed space group, predict
        // in P so the centering-absent reflections are integrated and the space-group search can
        // confirm or disprove centering (and catch a missed superstructure). A user-fixed space
        // group is trusted, so reject its absences here.
        .centering = experiment.GetGemmiSpaceGroup().has_value() ? outcome.symmetry.centering : 'P',
        .wedge_deg = std::fabs(wedge_deg),
        .mosaicity_deg = std::fabs(mos_deg),
        // FWHM -> sigma; 0 when monochromatic, leaving the prediction unchanged.
        .bandwidth_sigma = experiment.GetBandwidthFWHM().value_or(0.0f) / 2.3548f,
    };

    // Predict, then integrate with the selected integrator (box-sum or profile-fit).
    auto pred_start_time = std::chrono::steady_clock::now();
    auto nrefl = prediction.Calc(outcome.experiment, latt, settings_prediction);
    auto pred_end_time = std::chrono::steady_clock::now();
    msg.bragg_prediction_time_s = std::chrono::duration<float>(pred_end_time - pred_start_time).count();

    // The engine picks box-sum vs profile-fit internally from the experiment's IntegratorMode; the
    // caller's callback binds it to the right image (GPU-resident buffer, host buffer, or the assembled
    // FPGA image read straight on the CPU).
    auto integration_start_time = std::chrono::steady_clock::now();
    i_outcome.reflections = integrate(prediction.GetReflections(), nrefl, msg.number);
    msg.integrated_reflections = i_outcome.reflections.size();
    auto integration_end_time = std::chrono::steady_clock::now();
    msg.integration_time_s = std::chrono::duration<float>(integration_end_time - integration_start_time).count();

    constexpr size_t kMaxReflections = 10000;
    if (i_outcome.reflections.size() > kMaxReflections) {
        // Keep only smallest d (highest resolution)
        std::nth_element(i_outcome.reflections.begin(),
                         i_outcome.reflections.begin() + static_cast<long>(kMaxReflections),
                         i_outcome.reflections.end(),
                         [](const Reflection& a, const Reflection& b) {
                             return a.d < b.d;
                         });
        i_outcome.reflections.resize(kMaxReflections);

        // Optional: make output ordered by d (nice for downstream / debugging)
        std::sort(i_outcome.reflections.begin(), i_outcome.reflections.end(),
                  [](const Reflection& a, const Reflection& b) { return a.d < b.d; });
    }

    CalcISigma(msg, i_outcome.reflections);
    CalcWilsonBFactor(msg, i_outcome.reflections);

    ScaleImage(msg, i_outcome);

    // Copy reflections to outgoing message
    msg.reflections = i_outcome.reflections;

    // Persist the per-image result for the whole-run scaling/merge pass, unless the caller opted out
    // (viewer interactive use only needs the current image, returned above via msg).
    if (retain_outcomes_) {
        const std::unique_lock ul(reflections_mutex);
        integration_outcome[msg.number] = std::move(i_outcome);
    }
}

std::optional<IndexAndRefine::IndexingOutcome>
IndexAndRefine::DetermineRefineAnalyze(DataMessage &msg, const SpotFindingSettings &spot_finding_settings) {
    if (!indexer_ || !spot_finding_settings.indexing)
        return std::nullopt;

    IndexingOutcome outcome(experiment);

    if (rotation_indexer)
        outcome = DetermineLatticeAndSymmetryRotation(msg);
    else
        outcome = DetermineLatticeAndSymmetry(msg);

    if (!outcome.lattice_candidate)
        return std::nullopt;

    if (experiment.GetIndexingSettings().GetGeomRefinementAlgorithm() != GeomRefinementAlgorithmEnum::None)
        RefineGeometryIfNeeded(msg, outcome);

    if (!outcome.lattice_candidate.has_value())
        return std::nullopt;

    if (!AnalyzeIndexing(msg, outcome.experiment, *outcome.lattice_candidate, outcome.extra_lattice_candidates))
        return std::nullopt;

    {
        std::unique_lock ul(reflections_mutex);
        unit_cells[msg.number] = outcome.lattice_candidate->GetUnitCell();
    }
    msg.lattice_type = outcome.symmetry;
    return outcome;
}

void IndexAndRefine::ProcessImage(DataMessage &msg,
                                  const SpotFindingSettings &spot_finding_settings,
                                  BraggPrediction &prediction,
                                  const BraggIntegrateFn &integrate) {
    auto outcome = DetermineRefineAnalyze(msg, spot_finding_settings);
    if (outcome && spot_finding_settings.quick_integration)
        QuickPredictAndIntegrate(msg, spot_finding_settings, prediction, integrate, *outcome);
}

bool IndexAndRefine::IndexFrameOnly(DataMessage &msg, const SpotFindingSettings &spot_finding_settings) {
    return DetermineRefineAnalyze(msg, spot_finding_settings).has_value();
}

std::optional<RotationIndexerResult> IndexAndRefine::FinalizeRotationIndexing() {
    if (rotation_indexer) {
        if (const auto latt = rotation_indexer->GetLattice())
            return latt;

        rotation_indexer->RunIndexing();
        return rotation_indexer->GetLattice();
    }
    return {};
}

IndexAndRefine &IndexAndRefine::ReferenceIntensities(std::vector<MergedReflection> &reference) {
    // An external reference is trusted to be in the correct hand, so use it to break the merohedral
    // indexing ambiguity per image (serial stills index each crystal independently).
    reindex_resolver = std::make_unique<ReindexAmbiguityResolver>(experiment, reference);
    return *this;
}

void IndexAndRefine::ScaleImage(DataMessage &msg, IntegrationOutcome& outcome) {
    if (!reindex_resolver)
        return;

    // The external reference fixes the cell/space group, breaks the indexing ambiguity and reports CCref,
    // but is NEVER a scale anchor: scaling an image against a foreign dataset injects cross-dataset
    // systematics and is a worse reference than the data's own merge, so scaling self-references at the
    // post-measurement merge for both workflows. Rotation resolves the ambiguity globally and self-scales
    // in RotationScaleMerge (ChooseReindex / ReferenceIntensityCC), so there is nothing to do per image.
    // Stills resolve the merohedral ambiguity per image here (each crystal indexes in a random hand; pick
    // the hand best-correlated with the reference, once and for good).
    if (experiment.IsRotationIndexing())
        return;

    auto scaling_start_time = std::chrono::steady_clock::now();
    reindex_resolver->Resolve(outcome.reflections);
    auto scaling_end_time = std::chrono::steady_clock::now();
    msg.image_scale_time_s = std::chrono::duration<float>(scaling_end_time - scaling_start_time).count();
}

ScalingResult IndexAndRefine::ScaleAllImages(const std::vector<MergedReflection> &reference, size_t nthreads) {
    ScaleOnTheFly scaling(experiment, reference);
    scaling.Scale(integration_outcome, nthreads);
    scale_cc.resize(integration_outcome.size());

    for (int i = 0; i < integration_outcome.size(); i++)
        scale_cc.at(i) = integration_outcome[i].image_scale_cc.value_or(NAN);

    return ScalingResult(integration_outcome);
}

const std::vector<float> &IndexAndRefine::GetImageCC() const {
    return scale_cc;
}

const std::vector<std::optional<UnitCell> > & IndexAndRefine::GetUnitCells() const {
    return unit_cells;
}

std::optional<UnitCell> IndexAndRefine::GetConsensusUnitCell() const {
    const auto dist_tolerance = experiment.GetIndexingSettings().GetUnitCellDistTolerance();
    const auto angle_tolerance = experiment.GetIndexingSettings().GetUnitCellAngleTolerance_deg();

    if (rotation_indexer) {
        auto result = rotation_indexer->GetLattice();
        if (!result)
            return {};
        return result->lattice.GetUnitCell();
    }

    std::vector<UnitCell> cells;
    {
        std::unique_lock ul(reflections_mutex);
        cells.reserve(unit_cells.size());
        for (const auto &cell: unit_cells) {
            if (cell && cell->is_finite())
                cells.emplace_back(*cell);
        }
    }

    if (cells.empty())
        return {};

    if (experiment.GetUnitCell()) {
        std::vector<UnitCell> accepted;
        accepted.reserve(cells.size());

        for (const auto &cell: cells) {
            if (cell.is_close(*experiment.GetUnitCell(), dist_tolerance, angle_tolerance))
                accepted.emplace_back(cell);
        }

        return MeanUnitCell(accepted);
    }

    size_t best_count = 0;
    UnitCell best_reference{};

    for (const auto &ref: cells) {
        size_t count = 0;
        for (const auto &cell: cells) {
            if (cell.is_close(ref, dist_tolerance, angle_tolerance))
                ++count;
        }

        if (count > best_count) {
            best_count = count;
            best_reference = ref;
        }
    }

    if (best_count == 0)
        return {};

    std::vector<UnitCell> accepted;
    accepted.reserve(best_count);

    for (const auto &cell: cells) {
        if (cell.is_close(best_reference, dist_tolerance, angle_tolerance))
            accepted.emplace_back(cell);
    }

    return MeanUnitCell(accepted);
}

std::vector<IntegrationOutcome> &IndexAndRefine::GetIntegrationOutcome() {
    return integration_outcome;
}

const std::vector<IntegrationOutcome> &IndexAndRefine::GetIntegrationOutcome() const {
    return integration_outcome;
}

void IndexAndRefine::ForceRotationIndexerLattice(const CrystalLattice &lattice) {
    if (rotation_indexer)
        rotation_indexer->ForceLattice(lattice);
}

void IndexAndRefine::ForceRotationIndexerResult(const RotationIndexerResult &result) {
    if (rotation_indexer)
        rotation_indexer->ForceResult(result);
}