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

#include "../../common/ParallelFor.h"
#include "PostRefine.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <memory>
#include <numeric>

#include "../../common/JFJochMath.h"   // PI
#include "XtalResidual.h"              // XtalResidual (the positional detector<->reciprocal residual, step B)
#include "LatticeReduction.h"
#include "ceres/ceres.h"
#include "ceres/rotation.h"

namespace {

// One integrated partial, flattened across all images. Kept as narrow as the sort and the event
// split allow: on a large cell this array is gigabytes, and the scatter and every level of the
// per-bucket sort move all of it. The goniometer angle is not stored - it is a function of the
// image number alone, and is rebuilt from it where it is needed.
struct Partial {
    int h, k, l;
    float img;
    float I, sigma;
    float obs_x, obs_y;  // observed spot centroid (pixels); NAN if the box sum found no centroid
};

// A rocking event and its precomputed reference reciprocal vector (phi=0 frame, from the indexed lattice).
struct Event {
    double phi_obs;     // rad, intensity-weighted rocking centroid
    double weight;      // sqrt(sum I / sum sigma)
    double e_ref[3];    // h*a* + k*b* + l*c* at the reference (unrefined) cell/orientation
    int h, k, l;
};

// Distance-INDEPENDENT Ewald excitation residual for a uniform cell-scale parameter s and a refined
// goniometer axis (3-vector). The header-distance miscalibration leaves a uniform cell scale; the axis is
// the other phi_obs lever. Both are phi_obs-constrained (distance-independent). e_ref is the reference
// reciprocal (h*a* + k*b* + l*c* at the indexed cell). On the Ewald sphere <=> |p|^2 + 2 p_z/lambda == 0.
struct ScaleAxisExcitationResidual {
    ScaleAxisExcitationResidual(double lambda, double angle_rad, double weight, const double e_ref[3])
        : inv_lambda(1.0 / lambda), angle_rad(angle_rad), weight(weight),
          ex(e_ref[0]), ey(e_ref[1]), ez(e_ref[2]) {}
    template<typename T>
    bool operator()(const T *const s, const T *const axis, T *residual) const {
        const T inv_s = T(1) / s[0];
        const T p_ref[3] = {T(ex) * inv_s, T(ey) * inv_s, T(ez) * inv_s};
        const T aa[3] = {T(-angle_rad) * axis[0], T(-angle_rad) * axis[1], T(-angle_rad) * axis[2]};
        T p_lab[3];
        ceres::AngleAxisRotatePoint(aa, p_ref, p_lab);
        const T zeta = p_lab[0] * p_lab[0] + p_lab[1] * p_lab[1] + p_lab[2] * p_lab[2]
                       + T(2.0) * p_lab[2] * T(inv_lambda);
        residual[0] = T(weight) * zeta * T(0.5) / T(inv_lambda);
        return true;
    }
    const double inv_lambda, angle_rad, weight, ex, ey, ez;
};

// GONIOMETER ROTATION SCALE k: the same Ewald excitation residual, but with the cell scale and the axis
// DIRECTION already committed by step A, so the single free quantity is how far the stage actually turned
// per unit of commanded angle. Two differences from step A matter:
//   * the angle is measured from the CENTRE of the sweep, not from the goniometer's zero. The reference
//     orientation is the one rotation indexing fitted against the commanded angles, so it has already
//     absorbed the MEAN angle error; only the part that varies across the sweep is left to fit. Scaling the
//     absolute angle instead - which is what reading k off the length of step A's axis vector does - asks
//     the fit to also produce a constant offset it has no parameter for, and the least-squares compromise
//     shrinks k towards 1 by var(phi) / (var(phi) + phi_centre^2): exactly a factor of four for the common
//     case of a sweep starting at zero.
//   * e_mid is the reference reciprocal vector already turned to the sweep centre and divided by the
//     committed cell scale, so nothing but k is free.
struct RotationScaleResidual {
    RotationScaleResidual(double lambda, double dangle_rad, const double u[3], const double e_mid[3])
        : inv_lambda(1.0 / lambda), dangle_rad(dangle_rad),
          ux(u[0]), uy(u[1]), uz(u[2]), ex(e_mid[0]), ey(e_mid[1]), ez(e_mid[2]) {}
    template<typename T>
    bool operator()(const T *const k, T *residual) const {
        const T a = T(-dangle_rad) * k[0];
        const T aa[3] = {a * T(ux), a * T(uy), a * T(uz)};
        const T p_ref[3] = {T(ex), T(ey), T(ez)};
        T p_lab[3];
        ceres::AngleAxisRotatePoint(aa, p_ref, p_lab);
        const T zeta = p_lab[0] * p_lab[0] + p_lab[1] * p_lab[1] + p_lab[2] * p_lab[2]
                       + T(2.0) * p_lab[2] * T(inv_lambda);
        residual[0] = zeta * T(0.5) / T(inv_lambda);
        return true;
    }
    const double inv_lambda, dangle_rad, ux, uy, uz, ex, ey, ez;
};

}  // namespace

PostRefineResult PostRefineRotationGeometry(const std::vector<IntegrationOutcome> &outcomes,
                                            const GoniometerAxis &axis,
                                            const DiffractionGeometry &nominal_geom,
                                            const CrystalLattice &reference_latt,
                                            const PostRefineSettings &settings,
                                            Logger &logger) {
    PostRefineResult result;
    result.geom = nominal_geom;
    result.cell = reference_latt.GetUnitCell();
    result.distance_before_mm = nominal_geom.GetDetectorDistance_mm();
    result.distance_after_mm = nominal_geom.GetDetectorDistance_mm();
    try {
        const double wedge_half = axis.GetWedge_deg() / 2.0;
        const double lambda = nominal_geom.GetWavelength_A();
        const Coord ax = axis.GetAxis();
        const Coord Astar = reference_latt.Astar(), Bstar = reference_latt.Bstar(), Cstar = reference_latt.Cstar();

        // The goniometer angle is a function of the image number alone, so it is recomputed where it
        // is used rather than carried through the array below: eight bytes per partial cost more in
        // the fill, the scatter and every level of the sort than the multiply-add that rebuilds them.
        const auto angle_rad = [&](float img) { return (axis.GetAngle_deg(img) + wedge_half) * PI / 180.0; };

        // Count first, then fill. Growing one vector by push_back over tens of millions of
        // reflections copies the whole thing every time it doubles - several gigabytes of pure
        // copying - and the counts are cheap to take. Each outcome then owns a slice, so the fill
        // runs on all threads and lands in the order the serial loop produced. The h range comes out
        // of the same sweep: the sort below buckets by h and needs to know how many buckets that is,
        // and this pass already reads every reflection.
        const size_t nthreads = std::max(1, settings.num_threads);
        const int n_out = static_cast<int>(outcomes.size());
        std::vector<size_t> pts_offset(n_out + 1, 0);
        std::vector<int> h_lo_of(n_out), h_hi_of(n_out);
        ParallelChunks(n_out, nthreads, [&](int lo, int hi) {
            for (int o = lo; o < hi; o++) {
                size_t keep = 0;
                int lmin = std::numeric_limits<int>::max(), lmax = std::numeric_limits<int>::min();
                for (const auto &r : outcomes[o].reflections)
                    if (std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f) {
                        keep++;
                        lmin = std::min(lmin, r.h);
                        lmax = std::max(lmax, r.h);
                    }
                pts_offset[o + 1] = keep;
                h_lo_of[o] = lmin;
                h_hi_of[o] = lmax;
            }
        });
        int h_lo = std::numeric_limits<int>::max(), h_hi = std::numeric_limits<int>::min();
        for (int o = 0; o < n_out; o++) {
            pts_offset[o + 1] += pts_offset[o];
            h_lo = std::min(h_lo, h_lo_of[o]);
            h_hi = std::max(h_hi, h_hi_of[o]);
        }
        const size_t n_pts = pts_offset[n_out];
        const int H = (h_lo <= h_hi) ? (h_hi - h_lo + 1) : 1;

        // A vector of n partials VALUE-initialises them: on a large cell that is gigabytes of zeroing
        // on one thread, and it is that one thread which first touches every page - which on a
        // multi-socket machine leaves the whole array on its node, so every pass that follows runs at
        // one node's memory bandwidth. new[] leaves the partials untouched, so the parallel fill is
        // the first touch and each page lands on the node of the thread that filled it.
        std::unique_ptr<Partial[]> pts(new Partial[n_pts]);
        // The bucket histogram the sort needs is taken here rather than in a pass of its own, since
        // the fill already has h in hand. Its chunks are the outcome chunks ParallelChunks makes, so
        // the scatter below has to be split the same way.
        const int nt = static_cast<int>(std::clamp<size_t>(nthreads, 1, std::max(1, n_out)));
        const int chunk = (n_out + nt - 1) / nt;
        std::vector<std::vector<int32_t>> hist(nt, std::vector<int32_t>(H, 0));
        ParallelChunks(n_out, nthreads, [&](int lo, int hi) {
            std::vector<int32_t> &h_count = hist[lo / chunk];
            for (int o = lo; o < hi; o++) {
                size_t at = pts_offset[o];
                for (const auto &r : outcomes[o].reflections) {
                    if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f) continue;
                    const float ox = std::isfinite(r.observed_x) ? r.observed_x : NAN;
                    const float oy = std::isfinite(r.observed_y) ? r.observed_y : NAN;
                    pts[at++] = Partial{r.h, r.k, r.l, r.image_number, r.I, r.sigma, ox, oy};
                    h_count[r.h - h_lo]++;
                }
            }
        });
        logger.Info("Post-refine: {} partials gathered", n_pts);
        if (n_pts < static_cast<size_t>(settings.min_events)) return result;

        // Where each bucket starts, and the buckets largest first: the sort lays the array out this
        // way and the event split below walks the same buckets.
        std::vector<int32_t> bstart(H + 1, 0);
        std::vector<int> order(H);

        // Bucket by h, then sort the buckets. h is the leading key, so the sorted array is the
        // buckets laid end to end, and each bucket sorts on its own thread. Sorting the whole thing
        // in one pass moved every partial through every level of a comparison sort, on one thread,
        // over tens of millions of reflections.
        {
            // Not a total order: two partials of one reflection on one image still tie, as they did
            // before this was bucketed. What makes the result reproducible is the scatter below
            // rather than the comparator - the prefix lays each bucket out chunk by chunk and a
            // chunk is a contiguous span of the gathered order, so a bucket reaches std::sort in
            // global gather order whatever the thread count. Ties therefore resolve the same way on
            // every run and at every -N; they are simply not resolved by rank.
            const auto part_less = [](const Partial &a, const Partial &b) {
                if (a.h != b.h) return a.h < b.h;
                if (a.k != b.k) return a.k < b.k;
                if (a.l != b.l) return a.l < b.l;
                return a.img < b.img;
            };
            int32_t acc = 0;
            for (int b = 0; b < H; ++b) {
                bstart[b] = acc;
                for (int t = 0; t < nt; ++t) { const int32_t c = hist[t][b]; hist[t][b] = acc; acc += c; }
            }
            bstart[H] = acc;

            std::unique_ptr<Partial[]> sorted(new Partial[n_pts]);
            ParallelChunks(n_out, nthreads, [&](int lo, int hi) {
                std::vector<int32_t> fill = hist[lo / chunk];
                for (size_t i = pts_offset[lo]; i < pts_offset[hi]; ++i)
                    sorted[fill[pts[i].h - h_lo]++] = pts[i];
            });
            std::iota(order.begin(), order.end(), 0);
            std::sort(order.begin(), order.end(),
                      [&](int a, int b) { return (bstart[a + 1] - bstart[a]) > (bstart[b + 1] - bstart[b]); });
            ParallelFor(H, nthreads, [&](int oi) {
                const int b = order[oi];
                std::sort(sorted.get() + bstart[b], sorted.get() + bstart[b + 1], part_less);
            });
            pts.swap(sorted);
        }

        // Split into rocking events (same raw hkl, adjacent frames). Only >=2-frame events carry an
        // unbiased phi_obs (a single-frame centroid is just the frame centre); precompute e_ref per event.
        constexpr float MAX_FRAME_GAP = 2.0f;
        const auto run_end = [&](size_t i, size_t end) {
            size_t j = i + 1;
            while (j < end && pts[j].h == pts[i].h && pts[j].k == pts[i].k && pts[j].l == pts[i].l
                   && pts[j].img - pts[j - 1].img <= MAX_FRAME_GAP)
                ++j;
            return j;
        };
        // The event the partials [i, j) make, or false where their intensities cannot place a
        // centroid. Counting the events and writing them both walk the buckets, and both build the
        // event this way.
        const auto make_event = [&](size_t i, size_t j, Event &out) {
            double sumI = 0, sumIphi = 0, sumSig = 0;
            for (size_t m = i; m < j; ++m) {
                const double Ipos = std::max(0.0, static_cast<double>(pts[m].I));
                sumI += Ipos; sumIphi += Ipos * angle_rad(pts[m].img); sumSig += pts[m].sigma;
            }
            if (!(sumI > 0.0 && sumSig > 0.0)) return false;
            const Coord e = Astar * static_cast<float>(pts[i].h) + Bstar * static_cast<float>(pts[i].k)
                            + Cstar * static_cast<float>(pts[i].l);
            out = Event{sumIphi / sumI, std::sqrt(sumI / sumSig), {e.x, e.y, e.z},
                        pts[i].h, pts[i].k, pts[i].l};
            return true;
        };
        // An event never crosses an h boundary - h is the leading sort key - so the buckets can be
        // walked independently, and laying their events out in bucket order gives exactly the order
        // the serial walk produced. Counting first also sizes the array in one go, in place of a
        // push_back that grew a gigabyte by doubling.
        std::vector<int32_t> ev_count(H, 0);
        std::vector<size_t> ev_frames(H, 0);
        ParallelFor(H, nthreads, [&](int oi) {
            const int b = order[oi];
            const size_t end = bstart[b + 1];
            int c = 0;
            Event ev;
            for (size_t i = bstart[b]; i < end; ) {
                const size_t j = run_end(i, end);
                if (j - i >= 2 && make_event(i, j, ev)) ++c;
                i = j;
            }
            ev_count[b] = c;
        });
        std::vector<int32_t> ev_start(H + 1, 0);
        for (int b = 0; b < H; ++b) ev_start[b + 1] = ev_start[b] + ev_count[b];
        const size_t n_events = ev_start[H];
        std::unique_ptr<Event[]> events(new Event[n_events]);
        ParallelFor(H, nthreads, [&](int oi) {
            const int b = order[oi];
            const size_t end = bstart[b + 1];
            int at = ev_start[b];
            size_t frames = 0;
            for (size_t i = bstart[b]; i < end; ) {
                const size_t j = run_end(i, end);
                if (j - i >= 2 && make_event(i, j, events[at])) { frames += j - i; ++at; }
                i = j;
            }
            ev_frames[b] = frames;
        });
        size_t event_frames = 0;
        for (int b = 0; b < H; ++b) event_frames += ev_frames[b];
        // Frames per event is the phi_obs sampling: near 2 the reflections barely rock, so the angle
        // this refinement is fitted to is under-determined. It is a geometry count, so unlike an
        // intensity-weighted width it cannot be inflated by noise.
        logger.Info("Post-refine: {} multi-frame rocking events ({:.1f} frames per event)", n_events,
                    n_events == 0 ? 0.0 : static_cast<double>(event_frames) / n_events);
        if (static_cast<int>(n_events) < settings.min_events) return result;


        // The rotation-scale fit further down is a single scalar whose whole point is how the residual
        // varies ALONG the sweep, so it keeps every event. The cap below ranks by I/sigma, and on the
        // crystals that have a stage fault the strong events sit in the middle of the sweep - the part
        // that still indexes - so a capped set would leave the ends unrepresented in exactly the fit that
        // has to see them. Both sets come out of the one array by selecting on indices instead: with
        // the weights in the same places nth_element takes the same decisions it would take on the
        // events themselves, so the selection is the same one in the same order and the whole list no
        // longer has to be duplicated to survive it.
        constexpr size_t MAX_EVENTS = 20000;
        std::vector<int32_t> selected(n_events);
        std::iota(selected.begin(), selected.end(), 0);
        if (selected.size() > MAX_EVENTS) {
            std::nth_element(selected.begin(), selected.begin() + MAX_EVENTS, selected.end(),
                             [&](int32_t a, int32_t b) { return events[a].weight > events[b].weight; });
            selected.resize(MAX_EVENTS);
        }

        // ---- GEOMETRY REFINEMENT: the XtalOptimizer-equivalent, done as TWO SEPARATE
        // cross-validated steps rather than one joint fit (the same lesson as integration: refining the
        // profile width and the scale jointly fails, refining them separately works). Each step is committed
        // only if it lowers a HELD-OUT (deterministic split-half) residual - otherwise that part of the
        // geometry is left at nominal ("quit when things go wrong"):
        //   Step A: cell scale + rotation axis from phi_obs (distance-independent excitation residual).
        //   Step B: detector distance + beam centre from the observed spot positions, with the cell FIXED at
        //           step A (so the positional residual is no longer degenerate with the cell scale).
        // Detector tilt is held fixed (gauge-coupled to orientation on a single crystal). ----
        if (settings.refine_geometry) {
            const gemmi::CrystalSystem sys =
                (settings.crystal_system == gemmi::CrystalSystem::Trigonal) ? gemmi::CrystalSystem::Hexagonal
                                                                            : settings.crystal_system;
            const double ax0[3] = {ax.x, ax.y, ax.z};
            const double lambda_l = lambda;
            const double rot3 = nominal_geom.GetPoniRot3_rad();
            const double pixel_mm = nominal_geom.GetPixelSize_mm();
            const double det_rot[2] = {nominal_geom.GetPoniRot1_rad(), nominal_geom.GetPoniRot2_rad()};
            const UnitCell r0 = reference_latt.GetUnitCell();

            // Deterministic split of the reflections into a fit half and a held-out half. Avalanche-mix the
            // hkl hash so the split bit is decorrelated from the LSB - a plain h+k+l parity collides with the
            // lattice centering condition (e.g. an I-centred lattice has h+k+l even for EVERY present
            // reflection, so a parity split would leave the validation half empty).
            auto is_val = [](int h, int k, int l) {
                unsigned u = static_cast<unsigned>(h) * 2654435761u + static_cast<unsigned>(k) * 2246822519u
                             + static_cast<unsigned>(l) * 3266489917u;
                u ^= u >> 15; u *= 2246822519u; u ^= u >> 13;
                return (u & 1u) != 0u;
            };
            enum Subset { FIT, VAL, ALL };
            auto in = [&](int h, int k, int l, Subset s) {
                return s == ALL || (is_val(h, k, l) == (s == VAL)); };

            // ===== Step A: cell scale s + rotation axis from phi_obs =====
            auto excit_cost = [&](Subset s, double sc, const double axv[3]) {
                double c = 0.0; int n = 0;
                for (const int32_t i : selected) {
                    const Event &ev = events[i];
                    if (!in(ev.h, ev.k, ev.l, s)) continue;
                    ScaleAxisExcitationResidual r(lambda_l, ev.phi_obs, 1.0, ev.e_ref);
                    double sd = sc, av[3] = {axv[0], axv[1], axv[2]}, resid = 0.0;
                    r(&sd, av, &resid); c += resid * resid; ++n;
                }
                return n ? c / n : 0.0;
            };
            auto solve_scale_axis = [&](Subset s, double &s_out, double ax_out[3]) {
                double sc = 1.0, axv[3] = {ax0[0], ax0[1], ax0[2]};
                ceres::Problem p;
                for (const int32_t i : selected) {
                    const Event &ev = events[i];
                    if (!in(ev.h, ev.k, ev.l, s)) continue;
                    p.AddResidualBlock(new ceres::AutoDiffCostFunction<ScaleAxisExcitationResidual, 1, 1, 3>(
                        new ScaleAxisExcitationResidual(lambda_l, ev.phi_obs, settings.excitation_weight, ev.e_ref)),
                        new ceres::CauchyLoss(0.02), &sc, axv);
                }
                p.SetParameterLowerBound(&sc, 0, 0.9); p.SetParameterUpperBound(&sc, 0, 1.1);
                for (int j = 0; j < 3; ++j) { p.SetParameterLowerBound(axv, j, ax0[j] - 0.05);
                                              p.SetParameterUpperBound(axv, j, ax0[j] + 0.05); }
                ceres::Solver::Options o; o.linear_solver_type = ceres::DENSE_QR; o.max_num_iterations = 50;
                o.num_threads = std::max(1, settings.num_threads); o.logging_type = ceres::LoggingType::SILENT;
                ceres::Solver::Summary sum; ceres::Solve(o, &p, &sum);
                s_out = sc; ax_out[0] = axv[0]; ax_out[1] = axv[1]; ax_out[2] = axv[2];
                return sum.IsSolutionUsable();
            };
            double s_fit = 1.0, ax_fit[3];
            const bool convA = solve_scale_axis(FIT, s_fit, ax_fit);
            const double cvA_nom = excit_cost(VAL, 1.0, ax0);
            const double cvA_ref = excit_cost(VAL, s_fit, ax_fit);
            // Commit the cell scale only for a small, credible move: a well-calibrated header needs < ~0.6 %,
            // so a > 1 % scale is a red flag (on multi-lattice / noisy data the excitation fit is biased the
            // same way in every cross-validation fold, so the relative-improvement gate cannot catch it).
            result.cell_refined = convA && cvA_ref < 0.98 * cvA_nom && std::fabs(s_fit - 1.0) < 0.01;
            double s = 1.0, axv[3] = {ax0[0], ax0[1], ax0[2]};
            if (result.cell_refined) solve_scale_axis(ALL, s, axv);   // commit: re-fit on all data
            const double axlen = std::sqrt(axv[0]*axv[0] + axv[1]*axv[1] + axv[2]*axv[2]);
            const double axdev = std::acos(std::clamp((axv[0]*ax0[0]+axv[1]*ax0[1]+axv[2]*ax0[2])
                                                      / std::max(1e-9, axlen), -1.0, 1.0)) * 180.0 / PI;
            logger.Info("Post-refine GEOM step A (cell/axis): s = {:.5f}, rot-axis {:.3f} deg, held-out excit "
                        "{:.3e} -> {:.3e} => {}", s, axdev, cvA_nom, cvA_ref,
                        result.cell_refined ? "COMMIT" : "reject (kept nominal cell)");

            // ===== Goniometer rotation SCALE k, its own one-parameter fit on the same rocking events =====
            // The angles stored in the file are the COMMANDED ones, so a stage that turned k times as far
            // is invisible in the header. Nothing else here can represent it: the cell scale, the axis
            // direction, the distance and the beam are all orthogonal to a rotation MAGNITUDE error. Fitted
            // after step A so the cell scale and the axis direction are fixed at their committed values and
            // k is the only free quantity.
            const double u[3] = {axv[0] / axlen, axv[1] / axlen, axv[2] / axlen};
            double phi_c = 0.0, phi_lo = events[0].phi_obs, phi_hi = events[0].phi_obs;
            for (size_t e = 0; e < n_events; ++e) {
                phi_c += events[e].phi_obs;
                phi_lo = std::min(phi_lo, events[e].phi_obs);
                phi_hi = std::max(phi_hi, events[e].phi_obs);
            }
            phi_c /= static_cast<double>(n_events);
            const double sweep_deg = (phi_hi - phi_lo) * 180.0 / PI;
            // The reference reciprocal vector turned to the sweep centre, at the committed cell scale. The
            // angle then enters the fit measured FROM that centre. A constant crystal missetting about the
            // spindle is k with a slope in phi, so measuring the angle from the goniometer's zero instead
            // lets a missetting leak into k with gain <phi>/<phi^2> - which depends only on where the sweep
            // happens to sit. On a short sweep starting near zero that gain is enormous: a 0.14 deg
            // missetting on a 10 deg wedge fakes 1.4 % of k. Referred to the sweep centre the leak is
            // identically zero at any width, and no parameter has to be added to get it.
            // The residual is closed-form in k, so this is a one-parameter minimisation rather than a
            // solver problem. A rotation preserves length, so |p_lab| = |e_mid| whatever k is, and only
            // the z component moves; Rodrigues gives it exactly:
            //
            //   r(k) = C + A cos(a k) - B sin(a k) = C + R cos(a k + psi)
            //   C = lambda |e|^2 / 2 + u_z (u.e),  A = e_z - u_z (u.e),  B = (u x e)_z,  a = phi_obs - phi_c
            //
            // which is the same function the residual functor computes, to the last bit. Handing 8 million
            // one-parameter residual blocks to Ceres instead cost tens of millions of allocations and a
            // dense factorisation per iteration, for a fit that a scan over a bounded interval settles.
            // Coefficients are computed in double and stored narrowed: their rounding perturbs the
            // minimiser by ~1e-10, and k is carried downstream as a float.
            struct ScaleTerm { float a, C, R, psi; };
            std::vector<ScaleTerm> terms(n_events);
            std::vector<int> fifth_of(n_events);
            const double aa_c[3] = {-phi_c * u[0], -phi_c * u[1], -phi_c * u[2]};
            ParallelChunks(static_cast<int>(n_events), nthreads, [&](int lo, int hi) {
                for (int e = lo; e < hi; ++e) {
                    const double p[3] = {events[e].e_ref[0] / s, events[e].e_ref[1] / s,
                                         events[e].e_ref[2] / s};
                    double em[3];
                    ceres::AngleAxisRotatePoint(aa_c, p, em);
                    const double ue = u[0] * em[0] + u[1] * em[1] + u[2] * em[2];
                    const double e2 = em[0] * em[0] + em[1] * em[1] + em[2] * em[2];
                    const double C = 0.5 * lambda_l * e2 + u[2] * ue;
                    const double A = em[2] - u[2] * ue;
                    const double B = u[0] * em[1] - u[1] * em[0];
                    terms[e] = ScaleTerm{static_cast<float>(events[e].phi_obs - phi_c),
                                         static_cast<float>(C), static_cast<float>(std::hypot(A, B)),
                                         static_cast<float>(std::atan2(B, A))};
                    fifth_of[e] = std::clamp(static_cast<int>(
                            5.0 * (events[e].phi_obs - phi_lo) / std::max(1e-9, phi_hi - phi_lo)), 0, 4);
                }
            });

            // Robust-loss scale from the scatter the events actually have: it varies by more than a decade
            // between datasets, so any fixed constant is either inert or throws away real data. Taken once,
            // over every event, so the all-data fit and every jackknife fold share it.
            const auto residual_at = [&](const ScaleTerm &t, double k) {
                return static_cast<double>(t.C)
                     + static_cast<double>(t.R) * std::cos(static_cast<double>(t.a) * k + t.psi);
            };
            double rms = 0.0;
            for (const auto &t : terms) {
                const double r = residual_at(t, 1.0);
                rms += r * r;
            }
            rms = std::sqrt(rms / static_cast<double>(terms.size()));
            const double huber_delta = std::max(1e-12, 2.0 * rms);
            const double huber_d2 = huber_delta * huber_delta;

            // Ceres minimises half the sum of the loss applied to the SQUARED residual, so that is what is
            // reproduced here. One pass yields the five per-fifth partial sums, which serve the all-data
            // fit and all five leave-a-fifth-out folds together.
            // Each chunk folds into its own slot and the slots are summed in chunk order, so the
            // sum does not depend on which worker finishes first: the same events always add up in
            // the same sequence, and the fit is reproducible run to run.
            const int n_terms = static_cast<int>(terms.size());
            const int cost_nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads,
                                                                             static_cast<size_t>(n_terms))));
            const int cost_chunk = (n_terms + cost_nt - 1) / cost_nt;
            const auto cost_by_fifth = [&](double k) {
                std::vector<std::array<double, 5>> per_chunk(cost_nt);
                ParallelChunks(n_terms, nthreads, [&](int lo, int hi) {
                    std::array<double, 5> acc{};
                    for (int e = lo; e < hi; ++e) {
                        const double r = residual_at(terms[e], k);
                        const double s2 = r * r;
                        acc[fifth_of[e]] += (s2 <= huber_d2) ? s2
                                                             : (2.0 * huber_delta * std::sqrt(s2) - huber_d2);
                    }
                    per_chunk[lo / cost_chunk] = acc;
                });
                std::array<double, 5> total{};
                for (const auto &acc : per_chunk)
                    for (int j = 0; j < 5; ++j) total[j] += acc[j];
                return total;
            };

            // Scan the interval Ceres was bounded to, then close in. No event's phase can move by more than
            // a fraction of a period over an interval this narrow, so the objective has no structure the
            // grid could step over; the refinement is only there to place the minimum precisely.
            constexpr int SCALE_GRID = 101;
            constexpr double SCALE_K_LO = 0.95, SCALE_K_HI = 1.05;
            std::vector<std::array<double, 5>> grid(SCALE_GRID);
            for (int g = 0; g < SCALE_GRID; ++g)
                grid[g] = cost_by_fifth(SCALE_K_LO + (SCALE_K_HI - SCALE_K_LO) * g / (SCALE_GRID - 1));
            auto solve_scale = [&](int drop_fifth) {
                const auto total = [&](const std::array<double, 5> &f) {
                    double t = 0.0;
                    for (int j = 0; j < 5; ++j)
                        if (j != drop_fifth) t += f[j];
                    return t;
                };
                int best = 0;
                for (int g = 1; g < SCALE_GRID; ++g)
                    if (total(grid[g]) < total(grid[best])) best = g;
                const double step = (SCALE_K_HI - SCALE_K_LO) / (SCALE_GRID - 1);
                double a = std::max(SCALE_K_LO, SCALE_K_LO + step * (best - 1));
                double b = std::min(SCALE_K_HI, SCALE_K_LO + step * (best + 1));
                // Golden section: the objective is smooth but its curvature jumps wherever an event
                // crosses the Huber knee, which a derivative method would have to cope with.
                constexpr double INV_PHI = 0.6180339887498949;
                double c = b - INV_PHI * (b - a), d = a + INV_PHI * (b - a);
                double fc = total(cost_by_fifth(c)), fd = total(cost_by_fifth(d));
                // The fit is narrowed to a float before it is applied (Rugnux.h prepass_rotation_scale_),
                // so bracketing it below that type's epsilon, 6e-8, only buys about ten more full
                // passes over the events for a digit that cannot survive being stored.
                while (b - a > 1e-7) {
                    if (fc < fd) { b = d; d = c; fd = fc; c = b - INV_PHI * (b - a); fc = total(cost_by_fifth(c)); }
                    else         { a = c; c = d; fc = fd; d = a + INV_PHI * (b - a); fd = total(cost_by_fifth(d)); }
                }
                return 0.5 * (a + b);
            };
            const double k_fit = solve_scale(-1);
            result.rotation_scale = k_fit;

            // ----- Whether to COMMIT it. A stage fault is rare - 36 of 37 rotation datasets sit at 1.0000
            // on a direct scan - and a 1 % angle correction applied to a healthy dataset would damage it
            // silently, so every test below has to pass.
            // Preconditions: below these the fit is reported but never acted on. Under ~30 deg of sweep k
            // entangles with the axis direction and 10-20 deg truncations of a perfect dataset wander by
            // +-0.6 %; a screening wedge must not trigger a correction.
            constexpr int MIN_SCALE_EVENTS = 5000;
            constexpr double MIN_SCALE_SWEEP_DEG = 30.0;
            // T1 significance: 0.5 % is 18 sigma on the between-dataset scatter of healthy stages
            // (robust sd 2.8e-4) and still 3.5x below the one measured fault.
            constexpr double ROTATION_SCALE_TOL = 0.005;
            // T2 relevance: the misorientation the error produces at each end of the sweep. A large k over
            // a short sweep moves nothing and is not worth correcting.
            constexpr double MIN_SCALE_END_ERROR_DEG = 0.5;
            // T3 uniformity: a stage error is a ramp present in EVERY part of the sweep, so dropping any
            // fifth of it must leave the same k. A second lattice that dominates ONE END of the sweep -
            // exactly what happens where the primary stops indexing - fakes a k indistinguishable from a
            // real fault on T1 and T2, and is the reason this test is not optional. It replaces the
            // hkl-hash split used elsewhere here, which cannot see it: both halves of that split sit at
            // the same angles, so anything structured in phi survives in both folds.
            constexpr double MIN_SCALE_JACKKNIFE_FRAC = 0.5;
            const double end_error_deg = std::fabs(k_fit - 1.0) * sweep_deg / 2.0;
            const bool enough_data = static_cast<int>(n_events) >= MIN_SCALE_EVENTS
                                     && sweep_deg >= MIN_SCALE_SWEEP_DEG;
            const bool big_enough = enough_data && std::fabs(k_fit - 1.0) >= ROTATION_SCALE_TOL
                                    && end_error_deg >= MIN_SCALE_END_ERROR_DEG;
            double jackknife = 1.0;
            if (big_enough)
                for (int f = 0; f < 5; ++f)
                    jackknife = std::min(jackknife, (solve_scale(f) - 1.0) / (k_fit - 1.0));
            result.rotation_scale_suspect = big_enough && jackknife >= MIN_SCALE_JACKKNIFE_FRAC;
            logger.Info("Post-refine rotation SCALE: k = {:.5f} over {:.0f} deg of sweep centred on {:.1f} "
                        "deg ({} events): end error {:.2f} deg, leave-a-fifth-out {:.2f} => {}",
                        k_fit, sweep_deg, phi_c * 180.0 / PI, n_events, end_error_deg, jackknife,
                        result.rotation_scale_suspect ? "COMMIT"
                        : !enough_data ? "report only (too little sweep or too few events)"
                        : "reject (kept the stored angles)");
            if (result.rotation_scale_suspect)
                logger.Warning("Goniometer rotation scale looks off by {:+.2f} % (fitted {:.5f}): the stage "
                               "appears to have turned {} than the angles stored in the file, which are the "
                               "COMMANDED values. This is a hardware calibration fault, not a data problem - "
                               "left uncorrected it inflates mosaicity, biases the cell and loses "
                               "high-resolution reflections",
                               100.0 * (k_fit - 1.0), k_fit, k_fit > 1.0 ? "further" : "less far");

            // Cell (scale s, shape fixed) as the XtalResidual parameter blocks p0/p1/p2, held CONSTANT in step B.
            double p0[3] = {0, 0, 0}, p1[3] = {0, 0, 0}, p2[3] = {0, 0, 0};
            double beta = r0.beta;
            switch (sys) {
                case gemmi::CrystalSystem::Tetragonal:
                    LatticeToRodriguesAndLengths_GS(reference_latt, p0, p1);
                    p1[0] = (p1[0] + p1[1]) / 2.0; break;
                case gemmi::CrystalSystem::Cubic:
                    LatticeToRodriguesAndLengths_GS(reference_latt, p0, p1);
                    p1[0] = (p1[0] + p1[1] + p1[2]) / 3.0; break;
                case gemmi::CrystalSystem::Hexagonal:
                    LatticeToRodriguesAndLengths_Hex(reference_latt, p0, p1); break;
                case gemmi::CrystalSystem::Monoclinic:
                    LatticeToRodriguesLengthsBeta_Mono(reference_latt, p0, p1, beta);
                    p2[0] = beta; break;
                case gemmi::CrystalSystem::Orthorhombic:
                    LatticeToRodriguesAndLengths_GS(reference_latt, p0, p1); break;
                default:
                    LatticeToRodriguesAndLengths_GS(reference_latt, p0, p1);
                    p2[0] = r0.alpha * PI / 180.0; p2[1] = r0.beta * PI / 180.0; p2[2] = r0.gamma * PI / 180.0; break;
            }
            for (int j = 0; j < 3; ++j) p1[j] *= s;                     // apply the committed cell scale
            double rot_vec[3] = {axv[0], axv[1], axv[2]};              // committed (or nominal) axis

            // ===== Step B: detector distance + beam from the observed positions, cell fixed =====
            std::vector<const Partial *> obs;
            for (size_t i = 0; i < n_pts; ++i)
                if (std::isfinite(pts[i].obs_x) && std::isfinite(pts[i].obs_y)) obs.push_back(&pts[i]);
            constexpr size_t MAX_OBS = 20000;
            if (obs.size() > MAX_OBS) {
                std::nth_element(obs.begin(), obs.begin() + MAX_OBS, obs.end(),
                                 [](const Partial *a, const Partial *b) {
                                     return a->I / std::max(1e-9, static_cast<double>(a->sigma))
                                            > b->I / std::max(1e-9, static_cast<double>(b->sigma)); });
                obs.resize(MAX_OBS);
            }
            result.obs_used = static_cast<int>(obs.size());
            const double beam_x0 = nominal_geom.GetBeamX_pxl(), beam_y0 = nominal_geom.GetBeamY_pxl();
            const double dist0 = nominal_geom.GetDetectorDistance_mm();
            auto pos_cost = [&](Subset s, const double beam[2], const double dist[1]) {
                double c = 0.0; int n = 0;
                for (const Partial *pp : obs) {
                    if (!in(pp->h, pp->k, pp->l, s)) continue;
                    XtalResidual r(pp->obs_x, pp->obs_y, lambda_l, pixel_mm, rot3, angle_rad(pp->img),
                                   pp->h, pp->k, pp->l, sys);
                    double resid[3] = {0, 0, 0};
                    r(beam, dist, det_rot, rot_vec, p0, p1, p2, resid);
                    c += resid[0]*resid[0] + resid[1]*resid[1] + resid[2]*resid[2]; ++n;
                }
                return n ? c / n : 0.0;
            };
            auto solve_detector = [&](Subset s, double beam_out[2], double &dist_out) {
                double beam[2] = {beam_x0, beam_y0}, dist[1] = {dist0};
                ceres::Problem p;
                for (const Partial *pp : obs) {
                    if (!in(pp->h, pp->k, pp->l, s)) continue;
                    p.AddResidualBlock(new ceres::AutoDiffCostFunction<XtalResidual, 3, 2, 1, 2, 3, 3, 3, 3>(
                        new XtalResidual(pp->obs_x, pp->obs_y, lambda_l, pixel_mm, rot3, angle_rad(pp->img),
                                         pp->h, pp->k, pp->l, sys)),
                        new ceres::CauchyLoss(0.02), beam, dist,
                        const_cast<double *>(det_rot), rot_vec, p0, p1, p2);
                    p.SetParameterBlockConstant(const_cast<double *>(det_rot));
                    p.SetParameterBlockConstant(rot_vec);
                    p.SetParameterBlockConstant(p0); p.SetParameterBlockConstant(p1); p.SetParameterBlockConstant(p2);
                }
                if (p.NumResidualBlocks() == 0) { beam_out[0] = beam_x0; beam_out[1] = beam_y0; dist_out = dist0; return false; }
                p.SetParameterLowerBound(dist, 0, dist0 * 0.95); p.SetParameterUpperBound(dist, 0, dist0 * 1.05);
                for (int j = 0; j < 2; ++j) { p.SetParameterLowerBound(beam, j, beam[j] - 15.0);
                                              p.SetParameterUpperBound(beam, j, beam[j] + 15.0); }
                ceres::Solver::Options o; o.linear_solver_type = ceres::DENSE_QR; o.max_num_iterations = 60;
                o.num_threads = std::max(1, settings.num_threads); o.logging_type = ceres::LoggingType::SILENT;
                ceres::Solver::Summary sum; ceres::Solve(o, &p, &sum);
                beam_out[0] = beam[0]; beam_out[1] = beam[1]; dist_out = dist[0];
                return sum.IsSolutionUsable();
            };
            double beam[2] = {beam_x0, beam_y0}, dist = dist0;
            if (obs.size() >= static_cast<size_t>(settings.min_events)) {
                double beam_fit[2], dist_fit;
                const bool convB = solve_detector(FIT, beam_fit, dist_fit);
                const double b_nom[2] = {beam_x0, beam_y0}, d_nom[1] = {dist0};
                const double b_ref[2] = {beam_fit[0], beam_fit[1]}, d_ref[1] = {dist_fit};
                const double cvB_nom = pos_cost(VAL, b_nom, d_nom);
                const double cvB_ref = pos_cost(VAL, b_ref, d_ref);
                // Commit the detector geometry only for a small, credible move: distance < 1 % (a calibrated
                // header needs < ~0.6 %). A larger move is the red flag for an unreliable fit - typically a
                // second lattice whose spots bias every cross-validation fold identically, so the relative
                // "it improved" gate is blind to it and pulls a spurious distance<->cell pair (the radial
                // degeneracy) far off. The absolute size of the move discriminates a genuine header correction
                // from that failure far better than the absolute residual, which real marginal (noisy / iced)
                // data shares with the multi-lattice case.
                const bool in_bounds = std::fabs(dist_fit - dist0) < 0.01 * dist0
                                       && std::hypot(beam_fit[0] - beam_x0, beam_fit[1] - beam_y0) < 15.0;
                result.detector_refined = convB && cvB_ref < 0.98 * cvB_nom && in_bounds;
                if (result.detector_refined) { double bo[2]; solve_detector(ALL, bo, dist); beam[0] = bo[0]; beam[1] = bo[1]; }
                logger.Info("Post-refine GEOM step B (distance/beam): dist {:.3f} -> {:.3f} mm, beam "
                            "({:.2f},{:.2f}) -> ({:.2f},{:.2f}), held-out pos {:.3e} -> {:.3e} => {}",
                            dist0, result.detector_refined ? dist : dist0, beam_x0, beam_y0,
                            result.detector_refined ? beam[0] : beam_x0, result.detector_refined ? beam[1] : beam_y0,
                            cvB_nom, cvB_ref, result.detector_refined ? "COMMIT" : "reject (kept nominal detector)");
            } else {
                logger.Info("Post-refine GEOM step B: only {} positional observations - skipped", obs.size());
            }

            // Assemble the committed geometry.
            UnitCell cellA = r0;
            if (result.cell_refined) { cellA.a = static_cast<float>(r0.a * s); cellA.b = static_cast<float>(r0.b * s);
                                       cellA.c = static_cast<float>(r0.c * s); }
            result.cell = cellA;
            result.distance_after_mm = dist;
            result.beam_x_before_px = beam_x0; result.beam_x_after_px = beam[0];
            result.beam_y_before_px = beam_y0; result.beam_y_after_px = beam[1];
            result.events_used = static_cast<int>(selected.size());
            result.ok = result.cell_refined || result.detector_refined;
            if (!result.ok)
                logger.Info("Post-refine GEOM: neither step passed cross-validation - geometry left at nominal");
            return result;
        }

        return result;   // refine_geometry is the only supported mode; nothing refined otherwise
    } catch (...) {
        result.ok = false;
        return result;
    }
}