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

#include "BeamCenterFromSpots.h"

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

#include <algorithm>
#include <cmath>
#include <numeric>
#include <tuple>

namespace {

// The vote. A mirrored pair's two coordinates sum to twice the beam's, so the true pairs pile up at
// that value on a broad pedestal of accidental ones; a running mean over VOTE_BACKGROUND_BINS
// measures the pedestal and takes it away. +-60 px about twice the current guess covers any header
// error worth correcting, and 0.2 px bins are fine enough to seed the refinement.
constexpr float VOTE_BIN_PXL = 0.2f;
constexpr float VOTE_REACH_PXL = 60.0f;
constexpr int VOTE_BACKGROUND_BINS = 61;
constexpr int MIN_VOTES = 50;

// A reciprocal-lattice ROW makes the vote a comb: pairs one lattice step out of register vote at
// 2*beam + n*p, and such a tooth is sometimes TALLER than the true one. So every tooth above this
// fraction of the tallest, and this far from a taller one, is refined, and they are told apart
// afterwards on evidence the height does not carry.
constexpr float CANDIDATE_MIN_HEIGHT = 0.30f;
constexpr float CANDIDATE_SEPARATION_PXL = 6.0f;
constexpr int MAX_CANDIDATES = 6;

// The Friedel mirror does not touch the other coordinate, so a genuine pair agrees in it: loosely
// while the centre is still unknown (the vote), tightly once it is (the intensity correlation).
constexpr float FRIEDEL_OTHER_TOL_PXL = 3.0f;
constexpr float FRIEDEL_MATCH_WIN_PXL = 1.2f;
constexpr float FRIEDEL_OTHER_COST_WEIGHT = 0.05f;
constexpr float FRIEDEL_TIGHT_PXL = 0.8f;
constexpr int MIN_FRIEDEL_MATCHES = 20;
constexpr int MIN_MATCHES_PER_PAIR = 5;

// Fewer independent estimates than this - frame pairs for the Friedel fit, frames for the crossing
// - and there is nothing to take a scatter of.
constexpr int MIN_INDEPENDENT_ESTIMATES = 3;

// A frame at phi can only pair with one at phi+180, so on a sweep of S degrees the pairs' first
// frames span S - 180 however many of them there are - and it is that span, not their number, that
// the second crossing lives on: a reflection's two crossings are an angle apart and the sweep has to
// contain it. Just past half a turn the estimator still forms its full complement of pairs and
// answers from them, and the answer is tens of pixels out on most crystals: at 185 deg of sweep six
// of six committed answers were 2.6-28 px wrong and at 190 deg five of seven, all of it on the
// crossing coordinate while the Friedel one stayed inside 0.4 px; by 200 deg not one of 65 was
// wrong. The fit cannot notice this about itself - every pair agrees with every other - so it is a
// condition on the sweep, tested before anything is fitted. Measured on complete sweeps restricted
// to a span W, the crossing fit as it stood before the timing test above commits gross answers at
// W = 10, 15 and 20 and is clean only from 25; with the timing test it is clean at every W. So this
// is a floor under a guard that is now carried elsewhere, and its value is the largest that costs
// nothing measured: the shortest sweep in the regression set, 199.8 deg, yields a span of 17.6 deg -
// a sweep of S degrees gives S - 180, less one oscillation because the paired sample stops one frame
// short of a partner it would have no room for, less the frames the shutter margin keeps back at
// each end.
constexpr float MIN_PAIR_SPAN_DEG = 17.0f;

// A false tooth pairs unrelated reflections, and unrelated reflections do not have equal structure
// factors - |F(h)| = |F(-h)| holds for a real Friedel pair and for nothing else - so the intensity
// correlation of the matched set separates the teeth. It is a statistic and needs a population:
// below WEAK_MATCH_COUNT tight matches it is noise, and the count of them is the better evidence.
constexpr int WEAK_MATCH_COUNT = 200;
constexpr int MIN_CORRELATION_MATCHES = 30;

// The second crossing pairs spots within the whole pool rather than between two known frames, so
// its match has to be unambiguous: exactly one spot within the radius, and the spot must sit far
// enough off the mirror line that the two crossings are genuinely different measurements.
constexpr float CROSSING_OTHER_TOL_PXL = 0.6f;
constexpr float CROSSING_MATCH_RADIUS_PXL = 1.0f;
constexpr float CROSSING_MIN_LEVER_PXL = 25.0f;

// The crossing's own tooth test. The two crossings are the SAME reflection, so the geometry fixes
// not only where the second one is but WHEN: writing m1 = (m x s0)^ and m3 = (m1 x m)^ - the beam
// projected perpendicular to the spindle - the Laue condition holds q.m and q.m3 fixed along the
// sweep and lets only q.m1 change sign, so the two crossings are 2*atan2(q.m1, q.m3) of sweep apart.
// That angle is read off ONE spot's position, with no cell, no orientation and no indexing, and a
// false pairing has no reason to obey it: over the sampled frames it does so at the accidental rate
// of one frame in the turn. Measured over 36 datasets: true pairs keep it to a median 0.64 deg
// (worst 1.79), and the cut enriches them 29x - which is what tells the teeth apart, since the
// crossing's structure-factor correlation does not (it separates 13 of 17 against the timing's 17).
constexpr float CROSSING_PHI_TOL_DEG = 3.0f;

// A crossing pair is ONE measurement - the partner search finds it from both ends - so these count
// pairs, not matches. The 60-frame pre-scan yields 40 to 2100 of them (median 370), because both
// crossings of a reflection have to fall on sampled frames, which is quadratic in the budget: 30
// frames gives a median of 96 and 20 frames a median of 41.
constexpr int MIN_CROSSING_PAIRS = 10;
constexpr int MIN_CROSSING_PAIRS_PER_FRAME = 2;

// A reflection is recorded over a few consecutive frames, so its two crossings count as one
// measurement unless the frames are this many oscillation widths apart.
constexpr float MIN_EVENT_SEPARATION_WEDGES = 4.0f;

constexpr int REFINE_ITERATIONS = 6;
constexpr float REFINE_CONVERGED_PXL = 0.002f;

// Where the answer is checked from. The vote reaches only +-VOTE_REACH/2 px about wherever the
// search starts, so a taller tooth just outside that is invisible from the header but plain from a
// start displaced towards it; and the crossing fit sees the Friedel fit's coordinate, so a start
// elsewhere tests the two against each other as well. Four starts half a window away either come
// back to the same answer - and then it is the crystal's, not the header's - or find a competing
// one, which the estimator has no means of telling apart from it.
constexpr float CONSISTENCY_START_PXL = 25.0f;

// Persistent artefacts - the beam-stop halo, a hot pixel, the edge of a mask - sit at the SAME
// place on every frame, so they vote at twice their own position with one vote per frame pair, and
// that beats the real peak. A reflection never does: its Friedel mate is mirrored, not coincident.
// So a position that recurs across frames is not diffraction and is dropped.
constexpr float PERSISTENT_RADIUS_PXL = 1.5f;
constexpr float PERSISTENT_FRAME_FRACTION = 0.05f;
constexpr int PERSISTENT_MIN_FRAMES = 4;

// How far the spindle is searched. Both reaches are a real beamline's worst case with room to
// spare: measured against XDS's refined rotation axis over 38 sweeps the largest azimuth is
// 5.4 mrad and the largest tip 9.2 mrad. The grid is coarse on purpose - it exists to pick the
// TOOTH, and a least-squares fit on that tooth's own members then places the two angles far more
// finely than any grid could. Its step has to keep the tooth visible, which means keeping the
// spread it leaves - 2*step*(detector half-height) for the azimuth, 2*step*D*(1/cos2theta - 1) for
// the tip - to about a pixel.
constexpr float SPINDLE_AZIMUTH_REACH_RAD = 0.010f;
constexpr float SPINDLE_TIP_REACH_RAD = 0.020f;
constexpr float SPINDLE_GRID_STEP_RAD = 0.001f;
constexpr float SPINDLE_PEAK_WINDOW_PXL = 1.5f;
constexpr int SPINDLE_REFINE_ITERATIONS = 4;

float Coordinate(const BeamCenterSpot &spot, int axis) {
    return axis == 0 ? spot.x : spot.y;
}

float Coordinate(const std::pair<float, float> &point, int axis) {
    return axis == 0 ? point.first : point.second;
}

float Median(std::vector<float> v) {
    if (v.empty())
        return NAN;
    const size_t half = v.size() / 2;
    std::nth_element(v.begin(), v.begin() + half, v.end());
    return v[half];
}

// The scatter of a set, as a median absolute deviation scaled to a standard deviation.
float RobustSpread(const std::vector<float> &v) {
    const float centre = Median(v);
    std::vector<float> deviation(v.size());
    for (size_t i = 0; i < v.size(); i++)
        deviation[i] = std::abs(v[i] - centre);
    return 1.4826f * Median(deviation);
}

float BeamCoordinate(const DiffractionGeometry &geom, int axis) {
    return axis == 0 ? geom.GetBeamX_pxl() : geom.GetBeamY_pxl();
}

void SetBeamCoordinate(DiffractionGeometry &geom, int axis, float value) {
    if (axis == 0)
        geom.BeamX_pxl(value);
    else
        geom.BeamY_pxl(value);
}

// The exact lab-space mirror. Kept as an object because the trial centre changes far less often
// than the spots it is applied to.
class LabMirror {
    const DiffractionGeometry &geom;
    const RotMatrix inverse_rotation;
    const int axis;
public:
    LabMirror(const DiffractionGeometry &geometry, int mirror_axis)
        : geom(geometry), inverse_rotation(geometry.GetPoniRotMatrix().transpose()), axis(mirror_axis) {}

    // The image of a detector point: negate the mirrored lab component of the ray to it and project
    // the result back onto the detector. The reflecting plane contains the beam, so a point's image
    // sits opposite it about the DIRECT BEAM - not about the PONI, which the detector rotations put
    // up to (distance/pixel)*rot pixels away.
    [[nodiscard]] std::pair<float, float> operator()(float x, float y) const {
        Coord lab = geom.LabCoord(x, y);
        lab[axis] = -lab[axis];
        const Coord ray = inverse_rotation * lab;
        if (!(ray.z > 0))
            return {NAN, NAN};
        const float scale = geom.GetDetectorDistance_mm() / (ray.z * geom.GetPixelSize_mm());
        return {geom.GetBeamX_pxl() + ray.x * scale, geom.GetBeamY_pxl() + ray.y * scale};
    }
};

// The distance from the PONI to the direct beam on one axis. The vote is over detector coordinates
// mirrored about the direct beam, and the fit moves the PONI, so this is what converts one to the
// other. It does not depend on where the beam is, only on the detector rotations.
float DirectBeamOffset(const DiffractionGeometry &geom, int axis) {
    const auto direct = geom.GetDirectBeam_pxl();
    return Coordinate(direct, axis) - BeamCoordinate(geom, axis);
}

// The vote histogram with its pedestal taken away, bin by bin.
std::vector<float> VoteExcess(const std::vector<float> &sums, float guess) {
    const int n_bins = static_cast<int>(2 * VOTE_REACH_PXL / VOTE_BIN_PXL);
    const float first_bin = 2 * guess - VOTE_REACH_PXL;
    std::vector<float> histogram(n_bins, 0.0f);
    for (const float sum: sums) {
        // Floor, not truncation: a sum just below the first bin divides to a small negative number,
        // which truncates to zero and would pile every underflow into bin 0 - a tooth of its own,
        // and on real data a tall one.
        const int bin = static_cast<int>(std::floor((sum - first_bin) / VOTE_BIN_PXL));
        if (bin >= 0 && bin < n_bins)
            histogram[bin] += 1.0f;
    }

    std::vector<float> excess(n_bins);
    std::vector<double> cumulative(n_bins + 1, 0.0);
    for (int i = 0; i < n_bins; i++)
        cumulative[i + 1] = cumulative[i] + histogram[i];
    for (int i = 0; i < n_bins; i++) {
        const int lo = std::max(0, i - VOTE_BACKGROUND_BINS / 2);
        const int hi = std::min(n_bins, i + VOTE_BACKGROUND_BINS / 2 + 1);
        excess[i] = histogram[i] - static_cast<float>((cumulative[hi] - cumulative[lo]) / (hi - lo));
    }
    return excess;
}

// The values of 2*beam that the vote supports, tallest first: every bin whose excess over the
// pedestal is within CANDIDATE_MIN_HEIGHT of the tallest one's, and CANDIDATE_SEPARATION_PXL away
// from an already accepted candidate.
std::vector<float> VoteCandidates(const std::vector<float> &sums, float guess) {
    const int n_bins = static_cast<int>(2 * VOTE_REACH_PXL / VOTE_BIN_PXL);
    const float first_bin = 2 * guess - VOTE_REACH_PXL;
    const std::vector<float> excess = VoteExcess(sums, guess);

    std::vector<int> order(n_bins);
    std::iota(order.begin(), order.end(), 0);
    std::ranges::sort(order, [&](int a, int b) { return excess[a] > excess[b]; });

    std::vector<float> candidates;
    if (!(excess[order.front()] > 0.0f))
        return candidates;
    const float threshold = CANDIDATE_MIN_HEIGHT * excess[order.front()];
    for (const int bin: order) {
        if (excess[bin] < threshold || static_cast<int>(candidates.size()) >= MAX_CANDIDATES)
            break;
        const float value = first_bin + (bin + 0.5f) * VOTE_BIN_PXL;
        if (std::ranges::none_of(candidates, [&](float taken) {
                return std::abs(value - taken) < CANDIDATE_SEPARATION_PXL; }))
            candidates.push_back(value);
    }
    return candidates;
}

// One matched Friedel pair, in the frame where the mirror has already been applied.
struct FriedelMatch {
    float residual;      // the partner minus the mirror image, in the coordinate the mirror flips
    float other;         // and in the coordinate it leaves alone
    float log_intensity_a;
    float log_intensity_b;
};

// Mutual nearest neighbours between one frame's spots, mirrored, and its partner's. A spot may be
// matched only if it is its partner's best candidate as well, so a dense frame cannot pile several
// spots onto one.
void MatchFriedelFrames(const LabMirror &mirror, int axis,
                        const std::vector<BeamCenterSpot> &a, const std::vector<BeamCenterSpot> &b,
                        std::vector<FriedelMatch> &out) {
    const int other = 1 - axis;
    std::vector<int> best_of_a(a.size(), -1), best_of_b(b.size(), -1);
    std::vector<float> cost_of_a(a.size(), INFINITY), cost_of_b(b.size(), INFINITY);

    for (size_t i = 0; i < a.size(); i++) {
        const auto image = mirror(a[i].x, a[i].y);
        if (!std::isfinite(image.first))
            continue;
        for (size_t j = 0; j < b.size(); j++) {
            const float along = Coordinate(image, axis) - Coordinate(b[j], axis);
            const float across = Coordinate(image, other) - Coordinate(b[j], other);
            if (std::abs(along) >= FRIEDEL_MATCH_WIN_PXL || std::abs(across) >= FRIEDEL_OTHER_TOL_PXL)
                continue;
            const float cost = std::abs(along) + FRIEDEL_OTHER_COST_WEIGHT * std::abs(across);
            if (cost < cost_of_a[i]) { cost_of_a[i] = cost; best_of_a[i] = static_cast<int>(j); }
            if (cost < cost_of_b[j]) { cost_of_b[j] = cost; best_of_b[j] = static_cast<int>(i); }
        }
    }

    for (size_t i = 0; i < a.size(); i++) {
        const int j = best_of_a[i];
        if (j < 0 || best_of_b[j] != static_cast<int>(i))
            continue;
        const auto image = mirror(a[i].x, a[i].y);
        out.push_back({Coordinate(b[j], axis) - Coordinate(image, axis),
                       Coordinate(b[j], other) - Coordinate(image, other),
                       std::log(std::max(a[i].intensity, 1.0f)),
                       std::log(std::max(b[j].intensity, 1.0f))});
    }
}

struct FriedelCandidate {
    float poni = NAN;
    float sigma = NAN;
    int matches = 0;
    int tight_matches = 0;
    float correlation = -1.0f;
};

// Refine one tooth of the vote: mirror every frame onto its partner, and move the centre by half
// the median residual until it stops moving. Half, because a centre that is off by e puts the
// mirror image 2e away from the spot it belongs to.
std::optional<FriedelCandidate> RefineFriedel(DiffractionGeometry geom, int axis,
                                              const std::vector<std::vector<BeamCenterSpot>> &by_frame,
                                              const std::vector<std::pair<int, int>> &pairs,
                                              float start_direct_beam) {
    SetBeamCoordinate(geom, axis, start_direct_beam - DirectBeamOffset(geom, axis));

    for (int iteration = 0; iteration < REFINE_ITERATIONS; iteration++) {
        const LabMirror mirror(geom, axis);
        std::vector<FriedelMatch> matches;
        for (const auto &[first, second]: pairs)
            MatchFriedelFrames(mirror, axis, by_frame[first], by_frame[second], matches);
        if (static_cast<int>(matches.size()) < MIN_FRIEDEL_MATCHES)
            return {};

        std::vector<float> residual(matches.size());
        for (size_t i = 0; i < matches.size(); i++)
            residual[i] = matches[i].residual;
        const float step = Median(residual) / 2.0f;
        SetBeamCoordinate(geom, axis, BeamCoordinate(geom, axis) + step);
        if (std::abs(step) < REFINE_CONVERGED_PXL)
            break;
    }

    // The converged centre, and what each frame pair says about it on its own. The scatter of those
    // is the honest uncertainty: it is what a different pair of frames would have said, not how
    // finely their common median is determined, which is smaller by the square root of their number
    // and is a precision rather than an accuracy.
    const LabMirror mirror(geom, axis);
    std::vector<FriedelMatch> matches;
    std::vector<float> per_pair;
    for (const auto &[first, second]: pairs) {
        const size_t before = matches.size();
        MatchFriedelFrames(mirror, axis, by_frame[first], by_frame[second], matches);
        if (static_cast<int>(matches.size() - before) < MIN_MATCHES_PER_PAIR)
            continue;
        std::vector<float> residual;
        for (size_t i = before; i < matches.size(); i++)
            residual.push_back(matches[i].residual);
        per_pair.push_back(Median(residual) / 2.0f);
    }
    if (static_cast<int>(matches.size()) < MIN_FRIEDEL_MATCHES
        || static_cast<int>(per_pair.size()) < MIN_INDEPENDENT_ESTIMATES)
        return {};

    // A genuine Friedel pair sits at a constant offset in the coordinate the mirror does not touch;
    // a lattice-shifted one does not, so the correlation is measured over the ones that do.
    std::vector<float> across(matches.size());
    for (size_t i = 0; i < matches.size(); i++)
        across[i] = matches[i].other;
    const float centre = Median(across);

    std::vector<float> log_a, log_b;
    for (const auto &m: matches)
        if (std::abs(m.other - centre) < FRIEDEL_TIGHT_PXL) {
            log_a.push_back(m.log_intensity_a);
            log_b.push_back(m.log_intensity_b);
        }

    FriedelCandidate candidate;
    candidate.poni = BeamCoordinate(geom, axis);
    candidate.sigma = RobustSpread(per_pair);
    candidate.matches = static_cast<int>(matches.size());
    candidate.tight_matches = static_cast<int>(log_a.size());
    if (candidate.tight_matches > MIN_CORRELATION_MATCHES) {
        const float mean_a = std::accumulate(log_a.begin(), log_a.end(), 0.0f) / log_a.size();
        const float mean_b = std::accumulate(log_b.begin(), log_b.end(), 0.0f) / log_b.size();
        double covariance = 0, variance_a = 0, variance_b = 0;
        for (size_t i = 0; i < log_a.size(); i++) {
            covariance += (log_a[i] - mean_a) * (log_b[i] - mean_b);
            variance_a += (log_a[i] - mean_a) * (log_a[i] - mean_a);
            variance_b += (log_b[i] - mean_b) * (log_b[i] - mean_b);
        }
        if (variance_a > 0 && variance_b > 0)
            candidate.correlation = static_cast<float>(covariance / std::sqrt(variance_a * variance_b));
    }
    return candidate;
}

std::optional<FriedelCandidate> FitFriedel(const DiffractionGeometry &geom, int axis,
                                           const std::vector<std::vector<BeamCenterSpot>> &by_frame,
                                           const std::vector<std::pair<int, int>> &pairs) {
    const int other = 1 - axis;
    std::vector<float> sums;
    for (const auto &[first, second]: pairs)
        for (const auto &a: by_frame[first])
            for (const auto &b: by_frame[second])
                if (std::abs(Coordinate(a, other) - Coordinate(b, other)) < FRIEDEL_OTHER_TOL_PXL)
                    sums.push_back(Coordinate(a, axis) + Coordinate(b, axis));
    if (static_cast<int>(sums.size()) < MIN_VOTES)
        return {};

    std::vector<FriedelCandidate> refined;
    for (const float candidate: VoteCandidates(sums, Coordinate(geom.GetDirectBeam_pxl(), axis)))
        if (const auto fit = RefineFriedel(geom, axis, by_frame, pairs, candidate / 2.0f))
            refined.push_back(*fit);
    if (refined.empty())
        return {};

    // Which tooth is the real one. Normally the structure-factor correlation says so; where there
    // are too few tight matches for it to mean anything it is noise, and the tooth that matched the
    // most spots is the best evidence available.
    const bool weak = std::ranges::max_element(refined, {}, &FriedelCandidate::tight_matches)
                              ->tight_matches < WEAK_MATCH_COUNT;
    if (weak)
        return *std::ranges::max_element(refined, {}, &FriedelCandidate::tight_matches);
    return *std::ranges::max_element(refined, {}, &FriedelCandidate::correlation);
}

// How far along the sweep a spot's second crossing lies, for every spot. Kabsch's frame: m2 is the
// spindle, m3 the beam with its spindle component taken out, m1 the third. Along the sweep q.m2 and
// |q| never change and the Laue condition fixes q.m3, so a reflection reaches the sphere exactly
// twice, at azimuths +-atan2(q.m1, q.m3) about m3, and the sweep between them is twice that.
std::vector<float> CrossingSeparation_deg(const DiffractionGeometry &geom, const Coord &spindle,
                                          const std::vector<BeamCenterSpot> &spots) {
    const Coord m2 = spindle.Normalize();
    const Coord m1 = (m2 % geom.GetScatteringVector()).Normalize();
    const Coord m3 = (m1 % m2).Normalize();

    std::vector<float> separation(spots.size());
    for (size_t i = 0; i < spots.size(); i++) {
        const Coord q = geom.DetectorToRecip(spots[i].x, spots[i].y);
        separation[i] = 2.0f * std::atan2(q * m1, q * m3) * 180.0f / static_cast<float>(PI);
    }
    return separation;
}

// One tooth of the crossing vote, refined, and how much of the pool holds it up.
struct CrossingCandidate {
    float poni = NAN;
    float sigma = NAN;
    int pairs = 0;
};

// The other coordinate, from the second crossing. Both crossings of a reflection are somewhere in
// the pool rather than on two frames known in advance, so the match is over the whole pool: the
// spots are indexed on the coordinate the mirror leaves alone, and a window on it holds every
// candidate partner.
std::optional<CrossingCandidate> FitCrossing(DiffractionGeometry geom, int axis,
                                             const Coord &spindle,
                                             const std::vector<BeamCenterSpot> &spots,
                                             const std::vector<float> &frame_angle_deg,
                                             float min_separation_deg) {
    const int other = 1 - axis;
    std::vector<int> order(spots.size());
    std::iota(order.begin(), order.end(), 0);
    std::ranges::sort(order, [&](int a, int b) {
        return Coordinate(spots[a], other) < Coordinate(spots[b], other); });
    std::vector<float> across(spots.size());
    for (size_t i = 0; i < order.size(); i++)
        across[i] = Coordinate(spots[order[i]], other);

    const std::vector<float> separation = CrossingSeparation_deg(geom, spindle, spots);

    // Two spots can be the two crossings of one reflection when they are far enough apart in the
    // sweep to be two events at all, and when the sweep between them is the one the first spot's
    // own position asks for.
    const auto crossing_pair = [&](int a, int b) {
        const float delta = frame_angle_deg[spots[b].frame] - frame_angle_deg[spots[a].frame];
        return std::abs(delta) > min_separation_deg
               && std::abs(std::remainder(delta - separation[a], 360.0f)) < CROSSING_PHI_TOL_DEG;
    };

    std::vector<float> sums;
    for (size_t i = 0; i < order.size(); i++) {
        const auto last = std::upper_bound(across.begin() + i + 1, across.end(),
                                           across[i] + CROSSING_OTHER_TOL_PXL);
        for (auto it = across.begin() + i + 1; it != last; ++it) {
            const int j = order[it - across.begin()];
            if (crossing_pair(order[i], j))
                sums.push_back(Coordinate(spots[order[i]], axis) + Coordinate(spots[j], axis));
        }
    }
    if (static_cast<int>(sums.size()) < MIN_VOTES)
        return {};

    // Refine one tooth: mirror every spot, take the partner the mirror lands on, and move the centre
    // by half the median residual until it stops moving.
    const auto refine = [&](float start_2beam) -> std::optional<CrossingCandidate> {
        DiffractionGeometry trial = geom;
        SetBeamCoordinate(trial, axis, start_2beam / 2.0f - DirectBeamOffset(trial, axis));

        std::vector<float> residual;
        std::vector<int> frame;
        for (int iteration = 0; iteration <= REFINE_ITERATIONS; iteration++) {
            const LabMirror mirror(trial, axis);
            residual.clear();
            frame.clear();
            for (size_t q = 0; q < spots.size(); q++) {
                const auto image = mirror(spots[q].x, spots[q].y);
                if (!std::isfinite(image.first))
                    continue;
                // Both crossings measure the same |q|, so a spot close to the mirror line is its own
                // image and says nothing about where that line is.
                if (std::abs(Coordinate(spots[q], axis) - Coordinate(image, axis)) / 2.0f < CROSSING_MIN_LEVER_PXL)
                    continue;
                const float key = Coordinate(image, other);
                const auto first = std::lower_bound(across.begin(), across.end(), key - CROSSING_MATCH_RADIUS_PXL);
                const auto last = std::upper_bound(across.begin(), across.end(), key + CROSSING_MATCH_RADIUS_PXL);
                int found = -1;
                int count = 0;
                for (auto it = first; it != last; ++it) {
                    const int j = order[it - across.begin()];
                    if (!crossing_pair(static_cast<int>(q), j))
                        continue;
                    if (std::hypot(Coordinate(spots[j], axis) - Coordinate(image, axis),
                                   Coordinate(spots[j], other) - key) > CROSSING_MATCH_RADIUS_PXL)
                        continue;
                    found = j;
                    count++;
                }
                // The mirror is its own inverse, so the pair is found from both ends and gives the
                // same residual twice. It is one measurement and is counted once.
                if (count != 1 || found < static_cast<int>(q))
                    continue;
                residual.push_back(Coordinate(spots[found], axis) - Coordinate(image, axis));
                frame.push_back(spots[q].frame);
            }
            if (static_cast<int>(residual.size()) < MIN_CROSSING_PAIRS)
                return {};
            if (iteration == REFINE_ITERATIONS)
                break;

            const float step = Median(residual) / 2.0f;
            SetBeamCoordinate(trial, axis, BeamCoordinate(trial, axis) + step);
            if (std::abs(step) < REFINE_CONVERGED_PXL)
                break;
        }

        // As for the Friedel fit: the scatter of what the individual frames say, not of their mean.
        std::vector<std::vector<float>> of_frame(frame_angle_deg.size());
        for (size_t i = 0; i < residual.size(); i++)
            of_frame[frame[i]].push_back(residual[i] / 2.0f);
        std::vector<float> per_frame;
        for (const auto &one: of_frame)
            if (static_cast<int>(one.size()) >= MIN_CROSSING_PAIRS_PER_FRAME)
                per_frame.push_back(Median(one));
        if (static_cast<int>(per_frame.size()) < MIN_INDEPENDENT_ESTIMATES)
            return {};
        return CrossingCandidate{BeamCoordinate(trial, axis), RobustSpread(per_frame),
                                 static_cast<int>(residual.size())};
    };

    // Which tooth is the real one. A false tooth pairs unrelated reflections, and unrelated
    // reflections do not keep the sweep angle their positions ask for, so the pairs that survive
    // the timing test are the evidence the height does not carry - the crossing's counterpart of
    // the Friedel side's |F(h)| = |F(-h)|.
    std::optional<CrossingCandidate> best;
    for (const float candidate: VoteCandidates(sums, Coordinate(geom.GetDirectBeam_pxl(), axis)))
        if (const auto fit = refine(candidate))
            if (!best || fit->pairs > best->pairs)
                best = fit;
    return best;
}

// Drop the spots that are not diffraction: a position that appears on frame after frame.
std::vector<BeamCenterSpot> DropPersistentSpots(const std::vector<BeamCenterSpot> &spots, int n_frames) {
    std::vector<int> order(spots.size());
    std::iota(order.begin(), order.end(), 0);
    std::ranges::sort(order, [&](int a, int b) { return spots[a].x < spots[b].x; });
    std::vector<float> sorted_x(spots.size());
    for (size_t i = 0; i < order.size(); i++)
        sorted_x[i] = spots[order[i]].x;

    const int limit = std::max(PERSISTENT_MIN_FRAMES,
                               static_cast<int>(PERSISTENT_FRAME_FRACTION * n_frames));
    std::vector<BeamCenterSpot> kept;
    std::vector<int> frames;
    for (size_t i = 0; i < order.size(); i++) {
        const BeamCenterSpot &spot = spots[order[i]];
        const auto first = std::lower_bound(sorted_x.begin(), sorted_x.end(), spot.x - PERSISTENT_RADIUS_PXL);
        const auto last = std::upper_bound(sorted_x.begin(), sorted_x.end(), spot.x + PERSISTENT_RADIUS_PXL);
        frames.clear();
        for (auto it = first; it != last; ++it) {
            const BeamCenterSpot &other = spots[order[it - sorted_x.begin()]];
            if (std::hypot(other.x - spot.x, other.y - spot.y) <= PERSISTENT_RADIUS_PXL)
                frames.push_back(other.frame);
        }
        std::ranges::sort(frames);
        if (std::unique(frames.begin(), frames.end()) - frames.begin() <= limit)
            kept.push_back(spot);
    }
    return kept;
}

// THE SPINDLE IS NEVER PERPENDICULAR TO THE BEAM, AND THE FILE NEVER SAYS SO
//
// Everything above assumes the spindle m satisfies m.s0 = 0 and lies on a lab axis. Neither holds
// at a beamline, and the header cannot be asked: every master of the regression set writes the axis
// as exactly (-1,0,0) or (0,1,0), while XDS's refined axis for the same 38 sweeps departs from it
// by up to 5.4 mrad in the plane and 9.2 mrad out of it. So the deviation has to be measured here
// or not at all. It splits into two components that behave completely differently.
//
// AZIMUTH - the spindle turned about the beam by `a`. Both mirror planes turn with it (the
// Friedel one is normal to the spindle, the crossing one contains the spindle and the beam), so a
// mirror taken about the nominal lab axis leaves, along the mirrored coordinate, a residual of
// exactly 2*a*(the other coordinate). It does not bias the fit - that residual is odd in the other
// coordinate, so its median over a symmetric set is zero - it SMEARS the vote, by up to +-3 px at
// 2 mrad across a 1500 px detector, until a neighbouring comb tooth outvotes the true one. On a
// synthetic sweep that is 9.1 px of centre error at 2 mrad and 17.7 px at 5 mrad, reported at a
// sigma of 0.09-0.14 - a silent, confident, wrong answer. This is what is fitted below.
//
// TILT OUT OF THE PLANE - the spindle tipped towards the beam by `e`. Then the Friedel mate is no
// longer in diffracting condition at phi+180: with q' = q - 2(q.m)m the Ewald residual is
// -2(q.m)(m.s0), and the mate diffracts at phi+180+dphi with dphi = 2(q.m)e / (q.(m x s0_hat)).
// Carrying that rotation through to the detector leaves, along the mirrored coordinate, exactly
//
//     2 * e * D_pxl * (1/cos(2theta) - 1)
//
// (verified against the forward geometry to 1e-4 px), a purely RADIAL term - and the crossing axis
// does not see it at all, because its mirror plane's normal is m x s0, which is independent of the
// component of m along the beam. Uncorrected it costs the Friedel coordinate one tip times the
// median of that lever over the matched pairs: on the regression set a median 0.024 px per mrad,
// worst 0.060, i.e. a median 0.03 px and a worst case of 0.24 px. It is fitted below but NOT
// applied; the note at the call site says why.
//
// One thing this must NOT do is mirror about the spindle itself. The plane normal to m does not
// contain the beam once e is non-zero, so mirroring a ray about it moves the direct beam bodily by
// 2*e*D_pxl - 4.8 px, i.e. a 2.4 px centre error, at e = 2 mrad and 2400 px of distance. Only the
// azimuth of the mirror plane follows the spindle; the plane itself always contains the beam. That
// is why the correction below is a rotation of the whole problem about the beam and nothing else.

// A point turned about a centre, used to bring the sweep into the frame where the spindle does lie
// on a lab axis - which is the frame the mirror and both votes are written for.
std::pair<float, float> Turn(float x, float y, float cx, float cy, float sin_a, float cos_a) {
    const float dx = x - cx, dy = y - cy;
    return {cx + dx * cos_a - dy * sin_a, cy + dx * sin_a + dy * cos_a};
}

// How far the tip moves the mirror image of a spot, per radian of tip - half the pair residual
// above, which is what a shift of the spot itself has to be for the two to cancel.
float TipLever_pxl(float radius_pxl, float distance_pxl) {
    return std::hypot(radius_pxl, distance_pxl) - distance_pxl;
}

// One candidate Friedel pair, reduced to what the vote needs at any spindle: the sum of the two
// positions, their difference, and the two spots' tip levers together.
struct SpindlePair {
    float sum_x, sum_y;
    float diff_x, diff_y;
    float tip_lever;
};

// The spindle, from the Friedel vote.
//
// Neither component moves the true tooth, both spread it - the azimuth over 2*a*(detector height),
// the tip over 2*e*D*(1/cos2theta - 1) - so the spindle the spots were taken at is the one that
// makes the vote tallest. That is how the tooth is found; where it is found is then a least-squares
// question, because on the tooth's own members the vote value is linear in both angles:
//
//     value = 2*beam_along + (azimuth error)*(the pair's sum ACROSS the mirror line)
//                          - (tip error)*(the pair's tip lever)
//
// and those two columns are well separated - the first is odd across the mirror line, the second is
// even and radial. What is NOT in the fit is the third column a wrong detector rot1 would need, the
// mirrored coordinate squared; it is separable from the tip in principle (the two correlate at 0.66
// over a real spot distribution, condition number 6.1) but it is not separated here, which is the
// whole reason the tip is reported rather than used.
//
// A grid rather than a descent, because the two are not separable when both are far out: a 9 mrad
// tip - which the regression set has - spreads the tooth enough on its own that a spurious azimuth
// sharpens it, and a descent that meets the azimuth first never leaves that minimum.
//
// The pairs are collected once with a window wide enough for the whole search - the vote's own
// window, opened by the largest displacement an azimuth in range can produce - and every trial is
// then one pass over that list.
std::pair<float, float> FitSpindle(const std::vector<std::vector<BeamCenterSpot>> &by_frame,
                                   const std::vector<std::pair<int, int>> &pairs,
                                   int axis, float guess_x, float guess_y, float distance_pxl,
                                   float &excess_at_fit, float &excess_at_nominal) {
    const int other = 1 - axis;
    std::vector<SpindlePair> candidates;
    for (const auto &[first, second]: pairs)
        for (const auto &a: by_frame[first])
            for (const auto &b: by_frame[second]) {
                const float delta_other = Coordinate(a, other) - Coordinate(b, other);
                const float delta_axis = Coordinate(a, axis) - Coordinate(b, axis);
                if (std::abs(delta_other)
                    >= FRIEDEL_OTHER_TOL_PXL + SPINDLE_AZIMUTH_REACH_RAD * std::abs(delta_axis))
                    continue;
                candidates.push_back(
                    {a.x + b.x, a.y + b.y, a.x - b.x, a.y - b.y,
                     TipLever_pxl(std::hypot(a.x - guess_x, a.y - guess_y), distance_pxl)
                     + TipLever_pxl(std::hypot(b.x - guess_x, b.y - guess_y), distance_pxl)});
            }

    const float guess = axis == 0 ? guess_x : guess_y;
    const float centre_x = 2 * guess_x, centre_y = 2 * guess_y;

    // One trial: the pairs that still look like Friedel pairs at this spindle, and what each of
    // them votes for. `along` is the vote value, `across` the sum along the mirror line - the
    // azimuth's column in the fit below.
    std::vector<float> along, across, lever;
    const auto evaluate = [&](float trial_azimuth, float trial_tip) {
        const float sin_a = std::sin(trial_azimuth), cos_a = std::cos(trial_azimuth);
        along.clear(); across.clear(); lever.clear();
        for (const auto &c: candidates) {
            const float difference = other == 0 ? c.diff_x * cos_a + c.diff_y * sin_a
                                                : -c.diff_x * sin_a + c.diff_y * cos_a;
            if (std::abs(difference) >= FRIEDEL_OTHER_TOL_PXL)
                continue;
            const auto turned = Turn(c.sum_x, c.sum_y, centre_x, centre_y, -sin_a, cos_a);
            along.push_back(Coordinate(turned, axis) - trial_tip * c.tip_lever);
            across.push_back(Coordinate(turned, other));
            lever.push_back(c.tip_lever);
        }
    };
    const auto excess_of = [&](float trial_azimuth, float trial_tip) {
        evaluate(trial_azimuth, trial_tip);
        if (static_cast<int>(along.size()) < MIN_VOTES)
            return -INFINITY;
        const std::vector<float> excess = VoteExcess(along, guess);
        return *std::ranges::max_element(excess);
    };

    excess_at_nominal = excess_of(0.0f, 0.0f);
    float azimuth = 0.0f, tip = 0.0f, best = excess_at_nominal;
    const int azimuth_steps = static_cast<int>(SPINDLE_AZIMUTH_REACH_RAD / SPINDLE_GRID_STEP_RAD);
    const int tip_steps = static_cast<int>(SPINDLE_TIP_REACH_RAD / SPINDLE_GRID_STEP_RAD);
    for (int i = -azimuth_steps; i <= azimuth_steps; i++)
        for (int j = -tip_steps; j <= tip_steps; j++) {
            const float trial_azimuth = i * SPINDLE_GRID_STEP_RAD;
            const float trial_tip = j * SPINDLE_GRID_STEP_RAD;
            const float excess = excess_of(trial_azimuth, trial_tip);
            if (excess > best) { best = excess; azimuth = trial_azimuth; tip = trial_tip; }
        }
    excess_at_fit = best;

    // The tooth is now known to a grid step; where it sits is a straight three-parameter fit over
    // its own members, iterated because moving the angles changes which pairs land on it.
    for (int iteration = 0; iteration < SPINDLE_REFINE_ITERATIONS; iteration++) {
        evaluate(azimuth, tip);
        if (static_cast<int>(along.size()) < MIN_VOTES)
            return {azimuth, tip};
        const float peak = Median(along);
        double n = 0, s_a = 0, s_l = 0, s_aa = 0, s_al = 0, s_ll = 0, s_y = 0, s_ay = 0, s_ly = 0;
        for (size_t i = 0; i < along.size(); i++) {
            const double y = along[i] - peak;
            if (std::abs(y) >= SPINDLE_PEAK_WINDOW_PXL)
                continue;
            const double a = across[i], l = lever[i];
            n += 1; s_a += a; s_l += l; s_aa += a * a; s_al += a * l; s_ll += l * l;
            s_y += y; s_ay += a * y; s_ly += l * y;
        }
        if (n < MIN_VOTES)
            return {azimuth, tip};
        // The 2x2 system in (azimuth error, tip error) after the constant is projected out.
        const double caa = s_aa - s_a * s_a / n, cal = s_al - s_a * s_l / n, cll = s_ll - s_l * s_l / n;
        const double cay = s_ay - s_a * s_y / n, cly = s_ly - s_l * s_y / n;
        const double determinant = caa * cll - cal * cal;
        if (!(std::abs(determinant) > 0))
            return {azimuth, tip};
        azimuth -= static_cast<float>((cay * cll - cly * cal) / determinant);
        tip += static_cast<float>((cly * caa - cay * cal) / determinant);
        if (std::abs(azimuth) > SPINDLE_AZIMUTH_REACH_RAD || std::abs(tip) > SPINDLE_TIP_REACH_RAD)
            return {0.0f, 0.0f};
    }
    return {azimuth, tip};
}

// One run of the two fits, from wherever `geom` says the beam is. With `spindle_estimate` non-null
// the spindle is fitted here as well, and the sweep turned into its frame before anything else -
// so a run started elsewhere re-fits it, and a spindle that depends on where the search began shows
// up in the answer's spread like everything else.
std::optional<BeamCenterEstimate>
Estimate(DiffractionGeometry geom, const GoniometerAxis &goniometer,
         const std::vector<float> &frame_angle_deg, const std::vector<BeamCenterSpot> &spots,
         SpindleEstimate *spindle_estimate) {
    // Which detector coordinate the Friedel mirror flips: the one the spindle lies along. The lab
    // frame's x and y are the detector's own, so the spindle's larger lab component names the axis
    // - a vertical spindle simply swaps the two estimators over.
    const Coord spindle = goniometer.GetAxis().Normalize();
    const int friedel_axis = std::abs(spindle.x) >= std::abs(spindle.y) ? 0 : 1;

    // Frames half a turn apart. The pairing has to be as exact as the frames allow: a mate that
    // sits half an oscillation away is a different reflection.
    const float wedge = std::max(std::abs(goniometer.GetWedge_deg()), 1e-3f);
    std::vector<std::pair<int, int>> pairs;
    for (size_t i = 0; i < frame_angle_deg.size(); i++)
        for (size_t j = i + 1; j < frame_angle_deg.size(); j++)
            if (std::abs(std::abs(frame_angle_deg[j] - frame_angle_deg[i]) - 180.0f) < wedge)
                pairs.emplace_back(i, j);
    if (static_cast<int>(pairs.size()) < MIN_INDEPENDENT_ESTIMATES)
        return {};

    // The earlier frame of each pair, taken by ANGLE and not by position: the frames arrive sorted
    // on the first pass, but a second pass that reads more of the sweep appends them, and an index
    // order that is not an angle order would make this span the whole turn and the test vacuous.
    float earliest = INFINITY, latest = -INFINITY;
    for (const auto &pair: pairs) {
        const float first = std::min(frame_angle_deg[pair.first], frame_angle_deg[pair.second]);
        earliest = std::min(earliest, first);
        latest = std::max(latest, first);
    }
    if (latest - earliest < MIN_PAIR_SPAN_DEG)
        return {};

    auto kept = DropPersistentSpots(spots, static_cast<int>(frame_angle_deg.size()));
    std::vector<std::vector<BeamCenterSpot>> by_frame(frame_angle_deg.size());
    for (const auto &spot: kept)
        by_frame[spot.frame].push_back(spot);

    // The spindle, and the sweep turned by its azimuth about the beam. In that frame the spindle
    // does lie on a lab axis, so everything below is the estimator as written, and the answer is
    // turned back out of it at the end. Turning the SPOTS rather than the mirror is what carries the
    // correction into the two VOTES as well - they work in raw detector coordinates, where a mirror
    // line that is not parallel to a pixel axis breaks them just as thoroughly as it breaks the
    // mirror.
    //
    // The TIP is fitted alongside and then NOT applied. It has to be in the fit - left out, the
    // 9 mrad the regression set has spreads the vote enough on its own that a spurious azimuth
    // sharpens it - but it is a nuisance parameter, not a correction, and for two reasons. Its
    // column is radial, and so is what a wrong detector rot1 leaves behind, so the fitted tip is
    // the spindle's plus about 0.9 of the header's tilt error and is not the spindle alone
    // (measured over 38 sweeps: slope 1.03 on XDS's refined axis, offset +4.0 mrad). And its
    // scatter, about 2.6 mrad, times the lever it acts on is larger than the 0.07 px it would
    // correct - applied, it moves the centre a median 0.08 px the wrong way.
    const auto pivot = geom.GetDirectBeam_pxl();
    const float distance_pxl = geom.GetDetectorDistance_mm() / geom.GetPixelSize_mm();
    float azimuth = 0.0f;
    if (spindle_estimate) {
        float excess = NAN, excess_nominal = NAN, tip = NAN;
        std::tie(azimuth, tip) = FitSpindle(by_frame, pairs, friedel_axis, pivot.first, pivot.second,
                                            distance_pxl, excess, excess_nominal);
        *spindle_estimate = SpindleEstimate{azimuth, tip, excess, excess_nominal};
    }
    if (azimuth != 0.0f) {
        const float sin_a = std::sin(azimuth), cos_a = std::cos(azimuth);
        for (auto &spot: kept)
            std::tie(spot.x, spot.y) = Turn(spot.x, spot.y, pivot.first, pivot.second, -sin_a, cos_a);
        for (auto &frame: by_frame)
            frame.clear();
        for (const auto &spot: kept)
            by_frame[spot.frame].push_back(spot);
    }

    const auto friedel = FitFriedel(geom, friedel_axis, by_frame, pairs);
    if (!friedel)
        return {};
    SetBeamCoordinate(geom, friedel_axis, friedel->poni);

    const auto crossing = FitCrossing(geom, 1 - friedel_axis, spindle, kept, frame_angle_deg,
                                      MIN_EVENT_SEPARATION_WEDGES * wedge);
    if (!crossing)
        return {};
    SetBeamCoordinate(geom, 1 - friedel_axis, crossing->poni);

    if (azimuth != 0.0f) {
        // Back out of the spindle frame. The turn is about the beam, so it acts on the DIRECT beam
        // and the PONI follows from it: the offset between the two is a property of the detector
        // rotations alone and does not move.
        const float sin_a = std::sin(azimuth), cos_a = std::cos(azimuth);
        const auto direct = geom.GetDirectBeam_pxl();
        const auto turned = Turn(direct.first, direct.second, pivot.first, pivot.second, sin_a, cos_a);
        geom.BeamX_pxl(turned.first - DirectBeamOffset(geom, 0));
        geom.BeamY_pxl(turned.second - DirectBeamOffset(geom, 1));
    }

    return BeamCenterEstimate{geom.GetBeamX_pxl(), geom.GetBeamY_pxl(),
                              std::max(friedel->sigma, crossing->sigma)};
}

}   // namespace

std::optional<BeamCenterEstimate>
FindBeamCenterFromSpotSymmetry(const DiffractionExperiment &experiment,
                               const std::vector<float> &frame_angle_deg,
                               const std::vector<BeamCenterSpot> &spots,
                               SpindleEstimate *spindle_estimate) {
    const auto goniometer = experiment.GetGoniometer();
    if (!goniometer || spots.empty())
        return {};

    const auto geom = experiment.GetDiffractionGeometry();
    auto estimate = Estimate(geom, *goniometer, frame_angle_deg, spots, spindle_estimate);
    if (!estimate)
        return {};

    // How far the answer moves when the search is started somewhere else. On a sound measurement it
    // does not move at all; where it does, the two answers are both consistent with the spots and
    // nothing here can say which is the crystal's, so the scatter of the frame pairs - which stays
    // small for either of them - is not the uncertainty and this is. The spindle is re-fitted from
    // each start, so a fit that depends on where the search began is part of what is reported.
    float spread = 0.0f;
    for (const auto &[dx, dy]: {std::pair<float, float>{CONSISTENCY_START_PXL, 0.0f},
                                {-CONSISTENCY_START_PXL, 0.0f},
                                {0.0f, CONSISTENCY_START_PXL},
                                {0.0f, -CONSISTENCY_START_PXL}}) {
        DiffractionGeometry from = geom;
        from.BeamX_pxl(estimate->beam_x_pxl + dx).BeamY_pxl(estimate->beam_y_pxl + dy);
        SpindleEstimate again;
        if (const auto other = Estimate(from, *goniometer, frame_angle_deg, spots,
                                        spindle_estimate ? &again : nullptr))
            spread = std::max(spread, std::hypot(other->beam_x_pxl - estimate->beam_x_pxl,
                                                 other->beam_y_pxl - estimate->beam_y_pxl));
    }
    estimate->sigma_pxl = std::max(estimate->sigma_pxl, spread);
    return estimate;
}