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

#include "GeometryRefiner.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <vector>

#include "../../common/JFJochMath.h"   // PI
#include "../../common/ParallelFor.h"
#include "LatticeReduction.h"
#include "XtalResidual.h"
#include "ceres/ceres.h"

namespace {

// Trigonal is described with the hexagonal cell + B, same as the per-image XtalOptimizer.
gemmi::CrystalSystem MapSystem(gemmi::CrystalSystem s) {
    return (s == gemmi::CrystalSystem::Trigonal) ? gemmi::CrystalSystem::Hexagonal : s;
}

// The cell angles (radians) that XtalResidual bakes into its B matrix for a given system. Only
// monoclinic (beta) and triclinic (all three) read p2; the rest use fixed symmetry angles. We keep the
// angles constant in v1, so these are set once from the input cell.
void InitCellAngles(gemmi::CrystalSystem system, const UnitCell &cell, double ang[3]) {
    const double d2r = PI / 180.0;
    switch (system) {
        case gemmi::CrystalSystem::Hexagonal:
            ang[0] = PI / 2.0; ang[1] = PI / 2.0; ang[2] = 2.0 * PI / 3.0;
            break;
        case gemmi::CrystalSystem::Monoclinic:
            // XtalResidual reads p2[0] as beta.
            ang[0] = cell.beta * d2r; ang[1] = PI / 2.0; ang[2] = PI / 2.0;
            break;
        case gemmi::CrystalSystem::Triclinic:
            ang[0] = cell.alpha * d2r; ang[1] = cell.beta * d2r; ang[2] = cell.gamma * d2r;
            break;
        default:  // orthorhombic / tetragonal / cubic
            ang[0] = PI / 2.0; ang[1] = PI / 2.0; ang[2] = PI / 2.0;
            break;
    }
}

// Effective (lengths, alpha,beta,gamma) that reproduce the lattice XtalResidual builds from (len, ang),
// so a frame's lattice can be reconstructed with AngleAxisAndCellToLattice for HKL re-assignment.
void EffectiveCell(gemmi::CrystalSystem system, const double len[3], const double ang[3],
                   double lengths[3], double &alpha, double &beta, double &gamma) {
    switch (system) {
        case gemmi::CrystalSystem::Hexagonal:
            lengths[0] = len[0]; lengths[1] = len[0]; lengths[2] = len[2];
            alpha = PI / 2.0; beta = PI / 2.0; gamma = 2.0 * PI / 3.0;
            break;
        case gemmi::CrystalSystem::Tetragonal:
            lengths[0] = len[0]; lengths[1] = len[0]; lengths[2] = len[2];
            alpha = PI / 2.0; beta = PI / 2.0; gamma = PI / 2.0;
            break;
        case gemmi::CrystalSystem::Cubic:
            lengths[0] = len[0]; lengths[1] = len[0]; lengths[2] = len[0];
            alpha = PI / 2.0; beta = PI / 2.0; gamma = PI / 2.0;
            break;
        case gemmi::CrystalSystem::Monoclinic:
            lengths[0] = len[0]; lengths[1] = len[1]; lengths[2] = len[2];
            alpha = PI / 2.0; beta = ang[0]; gamma = PI / 2.0;
            break;
        case gemmi::CrystalSystem::Triclinic:
            lengths[0] = len[0]; lengths[1] = len[1]; lengths[2] = len[2];
            alpha = ang[0]; beta = ang[1]; gamma = ang[2];
            break;
        default:  // orthorhombic
            lengths[0] = len[0]; lengths[1] = len[1]; lengths[2] = len[2];
            alpha = PI / 2.0; beta = PI / 2.0; gamma = PI / 2.0;
            break;
    }
}

// Orientation seed (angle-axis) for a frame's lattice, using the per-system decomposition the per-image
// optimizer uses so it matches XtalResidual's B convention.
void OrientationSeed(gemmi::CrystalSystem system, const CrystalLattice &latt, double rod[3]) {
    double scratch[3];
    if (system == gemmi::CrystalSystem::Hexagonal) {
        LatticeToRodriguesAndLengths_Hex(latt, rod, scratch);
    } else if (system == gemmi::CrystalSystem::Monoclinic) {
        double beta;
        LatticeToRodriguesLengthsBeta_Mono(latt, rod, scratch, beta);
    } else {
        LatticeToRodriguesAndLengths_GS(latt, rod, scratch);
    }
}

UnitCell BuildUnitCell(gemmi::CrystalSystem system, const double len[3], const double ang[3]) {
    const double r2d = 180.0 / PI;
    UnitCell uc{};
    switch (system) {
        case gemmi::CrystalSystem::Hexagonal:
            uc = {(float)len[0], (float)len[0], (float)len[2], 90.0f, 90.0f, 120.0f};
            break;
        case gemmi::CrystalSystem::Tetragonal:
            uc = {(float)len[0], (float)len[0], (float)len[2], 90.0f, 90.0f, 90.0f};
            break;
        case gemmi::CrystalSystem::Cubic:
            uc = {(float)len[0], (float)len[0], (float)len[0], 90.0f, 90.0f, 90.0f};
            break;
        case gemmi::CrystalSystem::Monoclinic:
            uc = {(float)len[0], (float)len[1], (float)len[2], 90.0f, (float)(ang[0] * r2d), 90.0f};
            break;
        case gemmi::CrystalSystem::Triclinic:
            uc = {(float)len[0], (float)len[1], (float)len[2],
                  (float)(ang[0] * r2d), (float)(ang[1] * r2d), (float)(ang[2] * r2d)};
            break;
        default:  // orthorhombic
            uc = {(float)len[0], (float)len[1], (float)len[2], 90.0f, 90.0f, 90.0f};
            break;
    }
    return uc;
}

// Anchors the cell lengths to the input cell. Penalises each length's deviation from its input value
// with a large weight, breaking the low-resolution distance<->cell-scale degeneracy (position ~
// distance / cell_scale) so distance is determined by the data with the cell pinned.
struct CellLengthRegularizer {
    double w0, w1, w2;
    double l0, l1, l2;
    template <typename T>
    bool operator()(const T *const len, T *r) const {
        r[0] = T(w0) * (len[0] - T(l0));
        r[1] = T(w1) * (len[1] - T(l1));
        r[2] = T(w2) * (len[2] - T(l2));
        return true;
    }
};

// Where a spot found at (x, y) on module m really is (ModuleShift).
void PlaceOnModule(const ModuleShift &m, const double move[3], double x, double y, double &px, double &py) {
    const double ux = x - m.x0, uy = y - m.y0;
    const double c = std::cos(move[2]), s = std::sin(move[2]);
    px = m.x0 + c * ux - s * uy + move[0];
    py = m.y0 + s * ux + c * uy + move[1];
}

// How a residual reads the beam block. Without the tilt refined it is the PONI, as everywhere else. With the
// tilt refined it is the DIRECT BEAM, the point the undeflected beam lands on, and the PONI is worked out from
// it: a tilt moves the whole pattern by D tan(tilt)/pixel, exactly as a moved PONI does, so in PONI terms
// the two trade tens of pixels for each other (and the beam's box cannot hold the trade), while the direct
// beam does not move with the tilt and the tilt is left with only what no beam can imitate - the keystone.
// In XtalResidual's frame the direct beam is where M (u, v) = (-D tan rot1, -D tan rot2 / cos rot1), with
// (u, v) its offset from the PONI in mm and M the image orientation (a mirror and quarter turns, so M^-1 = M^T).
struct BeamBlock {
    bool direct = false;
    double m00 = 1, m01 = 0, m10 = 0, m11 = 1;   // DetectorOrientation::Matrix(), as XtalResidual takes it
    double pixel = 1;

    template <typename T, typename R>
    void Poni(const T *const beam, const T &distance_mm, const R *const detector_rot, T *poni) const {
        if (!direct) {
            poni[0] = beam[0];
            poni[1] = beam[1];
            return;
        }
        using std::cos;
        using std::tan;
        const T dx = -distance_mm * tan(detector_rot[0]);
        const T dy = -distance_mm * tan(detector_rot[1]) / cos(detector_rot[0]);
        poni[0] = beam[0] - (T(m00) * dx + T(m10) * dy) / T(pixel);
        poni[1] = beam[1] - (T(m01) * dx + T(m11) * dy) / T(pixel);
    }
};

// The residual of one spot (XtalResidual) on the blocks this refinement can move, in this order: the beam
// (read through BeamBlock) and the distance as one block, the tilt where it is refined, the frame's orientation,
// the cell lengths, and the module's move (dx, dy, rot) where the spot's module is placed. The rotation axis, the
// cell angles and a held tilt enter as plain numbers.
//
// The residual is the difference of two halves that share no block: the observed reciprocal vector (beam,
// distance, tilt, move) and the predicted one (orientation, cell). Ceres' AutoDiffCostFunction would carry the
// derivatives of every block through both halves; here each half is differentiated on its own blocks only, which
// gives the same values and the same derivatives with half as many of them per number.
//
// A placed module: moving the spot by R(rot) u - u + (dx, dy), with (ux, uy) its offset from the module
// centre, shifts its offset from the beam exactly as moving the beam the other way does, so the beam handed
// to XtalResidual carries the move and everything else is the residual of an unmoved spot.
class SpotCost : public ceres::CostFunction {
public:
    SpotCost(const XtalResidual &residual, const BeamBlock &beam_block, const double *detector_rot,
             const double *rotation_axis, const double *cell_ang, bool tilt, bool module, double ux = 0,
             double uy = 0)
        : residual(residual), beam_block(beam_block), detector_rot(detector_rot), rotation_axis(rotation_axis),
          cell_ang(cell_ang), tilt(tilt), module(module), ux(ux), uy(uy) {
        set_num_residuals(3);
        std::vector<int32_t> &sizes = *mutable_parameter_block_sizes();
        sizes.push_back(3);
        if (tilt)
            sizes.push_back(2);
        sizes.push_back(3);
        sizes.push_back(3);
        if (module)
            sizes.push_back(3);
    }

    bool Evaluate(double const *const *parameters, double *residuals, double **jacobians) const override {
        if (tilt)
            return module ? EvaluateWith<true, true>(parameters, residuals, jacobians)
                          : EvaluateWith<true, false>(parameters, residuals, jacobians);
        return module ? EvaluateWith<false, true>(parameters, residuals, jacobians)
                      : EvaluateWith<false, false>(parameters, residuals, jacobians);
    }

private:
    template <bool TILT, bool MODULE>
    bool EvaluateWith(double const *const *p, double *out, double **jac) const {
        constexpr int N_OBS = 3 + (TILT ? 2 : 0) + (MODULE ? 3 : 0);
        using JO = ceres::Jet<double, N_OBS>;
        using JP = ceres::Jet<double, 6>;
        constexpr int TILT_BLOCK = 1, P0_BLOCK = TILT ? 2 : 1, P1_BLOCK = P0_BLOCK + 1, MOVE_BLOCK = P1_BLOCK + 1;

        // The observed half, on the beam and distance, the tilt and the move.
        JO bd[3], t[2], m[3];
        for (int j = 0; j < 3; ++j)
            bd[j] = JO(p[0][j], j);
        if constexpr (TILT)
            for (int j = 0; j < 2; ++j)
                t[j] = JO(p[TILT_BLOCK][j], 3 + j);
        if constexpr (MODULE)
            for (int j = 0; j < 3; ++j)
                m[j] = JO(p[MOVE_BLOCK][j], N_OBS - 3 + j);
        JO obs[3];
        if constexpr (TILT)
            Observed(bd, t, MODULE ? m : nullptr, obs);
        else
            Observed(bd, detector_rot, MODULE ? m : nullptr, obs);

        // The predicted half, on the orientation and the cell lengths.
        JP p0[3], p1[3];
        for (int j = 0; j < 3; ++j) {
            p0[j] = JP(p[P0_BLOCK][j], j);
            p1[j] = JP(p[P1_BLOCK][j], 3 + j);
        }
        const JP p2[3] = {JP(cell_ang[0]), JP(cell_ang[1]), JP(cell_ang[2])};
        JP recip_unrot[3];
        residual.PredictedRecipUnrot(p1, p2, recip_unrot);
        JP pred[3];
        const AngleAxisRotator<JP> rot_p0(p0);
        rot_p0.Rotate(recip_unrot, pred);

        for (int r = 0; r < 3; ++r)
            out[r] = obs[r].a - pred[r].a;
        if (!jac)
            return true;
        // Row-major 3 x size blocks; a block held constant has no Jacobian asked for.
        auto fill = [](double *block, int size, auto derivative) {
            if (block)
                for (int r = 0; r < 3; ++r)
                    for (int c = 0; c < size; ++c)
                        block[r * size + c] = derivative(r, c);
        };
        fill(jac[0], 3, [&](int r, int c) { return obs[r].v[c]; });
        if constexpr (TILT)
            fill(jac[TILT_BLOCK], 2, [&](int r, int c) { return obs[r].v[3 + c]; });
        if constexpr (MODULE)
            fill(jac[MOVE_BLOCK], 3, [&](int r, int c) { return obs[r].v[N_OBS - 3 + c]; });
        fill(jac[P0_BLOCK], 3, [&](int r, int c) { return -pred[r].v[c]; });
        fill(jac[P1_BLOCK], 3, [&](int r, int c) { return -pred[r].v[3 + c]; });
        return true;
    }

    template <typename T, typename R>
    void Observed(const T *const beam_distance, const R *const tilt, const T *const move, T *recip_obs) const {
        using std::cos;
        using std::sin;
        T b[2];
        beam_block.Poni(beam_distance, beam_distance[2], tilt, b);
        if (move) {
            const T c = cos(move[2]), s = sin(move[2]);
            b[0] = b[0] - ((c - T(1)) * T(ux) - s * T(uy) + move[0]);
            b[1] = b[1] - (s * T(ux) + (c - T(1)) * T(uy) + move[1]);
        }
        const R axis[3] = {R(rotation_axis[0]), R(rotation_axis[1]), R(rotation_axis[2])};
        residual.ObservedRecip(b, beam_distance[2], tilt, axis, recip_obs);
    }

    const XtalResidual residual;
    const BeamBlock beam_block;
    const double *detector_rot;   // the tilt where it is held
    const double *rotation_axis;
    const double *cell_ang;
    const bool tilt, module;
    const double ux, uy;
};

// The residual of a fit that moves only the crystal orientations (refine_beam_distance off, no module
// placed): the observed reciprocal vector and the unrotated prediction are then constants of the spot, taken
// once in plain numbers, and only the orientation's three derivatives are carried.
struct OrientationResidual {
    double recip_obs[3];
    double recip_unrot[3];

    template <typename T>
    bool operator()(const T *const p0, T *out) const {
        T e_pred_recip[3];
        const AngleAxisRotator<T> rot_p0(p0);
        rot_p0.Rotate(recip_unrot, e_pred_recip);
        out[0] = recip_obs[0] - e_pred_recip[0];
        out[1] = recip_obs[1] - e_pred_recip[1];
        out[2] = recip_obs[2] - e_pred_recip[2];
        return true;
    }
};

// What module moves cannot be told apart from: all modules shifted alike is the beam centre, all turned
// alike (about their own centres) is the crystal orientations, and all pushed outwards alike is the
// distance. The sums of those three over the moved modules are held at zero - the beam, the orientations
// and the distance then carry them - so the modules are left with what only they can explain. The sums are
// weighted by the spots each module holds: a module the frames barely reach is placed poorly, and an
// unweighted sum would hand its error to all the others.
//
// The detector tilt is not among them. Within one module a tilt is a gradient (a scale and a shear across
// it), which no rigid move of the module makes, and with the tilt refined the placement gains no new
// direction it cannot tell apart; holding the modules' moves to no tilt-like pattern as well only cost
// accuracy on frames the fit had not seen.
struct ModuleGauge {
    std::vector<std::array<double, 2>> outward;   // unit vector from the beam to each module centre
    std::vector<double> share;                    // each module's fraction of the spots
    double lever_px;                              // turns are weighed at this distance from the centre
    double weight;

    template <typename T>
    bool operator()(T const *const *move, T *out) const {
        T sx(0), sy(0), sr(0), so(0);
        for (size_t i = 0; i < outward.size(); ++i) {
            const T w(share[i]);
            sx += w * move[i][0];
            sy += w * move[i][1];
            sr += w * move[i][2] * T(lever_px);
            so += w * (move[i][0] * T(outward[i][0]) + move[i][1] * T(outward[i][1]));
        }
        out[0] = T(weight) * sx;
        out[1] = T(weight) * sy;
        out[2] = T(weight) * sr;
        out[3] = T(weight) * so;
        return true;
    }
};

}   // namespace

GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_geom,
                                           const UnitCell &input_cell,
                                           const std::vector<GeomRefineFrame> &frames,
                                           const GeometryRefinerSettings &settings) {
    GeometryRefinerResult result;
    try {
        const gemmi::CrystalSystem system = MapSystem(settings.crystal_system);

        // Shared geometry parameter blocks (refined). The rotation axis is constant, and so is the detector tilt
        // unless settings.refine_tilt. The beam and the distance are one block of three: every block refined
        // beside it has three parameters too (orientations, cell lengths, module moves), so Ceres eliminates the
        // orientations with its fixed-size Schur code rather than the general one (1.7x faster on one thread).
        double beam_distance[3] = {nominal_geom.GetBeamX_pxl(), nominal_geom.GetBeamY_pxl(),
                                   nominal_geom.GetDetectorDistance_mm()};
        double *const beam = beam_distance;
        double &distance_mm = beam_distance[2];
        double detector_rot[2] = {nominal_geom.GetPoniRot1_rad(), nominal_geom.GetPoniRot2_rad()};
        double rot_vec[3] = {1.0, 0.0, 0.0};
        const double rot3 = nominal_geom.GetPoniRot3_rad();
        // Same for every spot of every frame, so taken once here rather than per residual.
        const double cos_rot3 = std::cos(rot3), sin_rot3 = std::sin(rot3);
        const DetectorOrientation orientation = nominal_geom.GetOrientation();
        const double lambda = nominal_geom.GetWavelength_A();
        const double pixel = nominal_geom.GetPixelSize_mm();

        // With the tilt refined, the beam block is the direct beam (BeamBlock), started and bounded where the
        // nominal geometry puts the direct beam.
        const bool refine_tilt = settings.refine_tilt && settings.refine_beam_distance;
        BeamBlock beam_block;
        if (refine_tilt) {
            const RotMatrix m = orientation.Matrix();
            beam_block = {true, m.Column(0).x, m.Column(1).x, m.Column(0).y, m.Column(1).y, pixel};
            double poni_at_zero[2];
            const double zero[2] = {0, 0};
            beam_block.Poni(zero, distance_mm, detector_rot, poni_at_zero);
            beam[0] -= poni_at_zero[0];
            beam[1] -= poni_at_zero[1];
        }
        const double beam0[2] = {beam[0], beam[1]};
        // The PONI the beam block stands for, at the current tilt and distance.
        auto poni = [&](double &x, double &y) {
            double p[2];
            beam_block.Poni(beam, distance_mm, detector_rot, p);
            x = p[0];
            y = p[1];
        };

        // Shared cell blocks. Lengths refined (regularized); angles constant.
        double cell_len[3] = {input_cell.a, input_cell.b, input_cell.c};
        const double cell_len0[3] = {input_cell.a, input_cell.b, input_cell.c};
        double cell_ang[3];
        InitCellAngles(system, input_cell, cell_ang);

        double eff_len[3], eff_alpha, eff_beta, eff_gamma;   // recomputed each round from cell_len/ang

        // Per-frame orientation seeds.
        std::vector<std::array<double, 3>> orient(frames.size());
        for (size_t i = 0; i < frames.size(); ++i)
            OrientationSeed(system, frames[i].lattice, orient[i].data());

        // Reciprocal length of one detector pixel at the nominal geometry, used to scale the robust
        // loss (~3 px) into the reciprocal-space residual units.
        const double recip_per_px = pixel / (distance_mm * lambda);
        const double loss_scale = 3.0 * recip_per_px;

        // The modules' moves (dx, dy, rot), starting where settings put them. Which of them are refined is
        // decided once, on the indexed spots each holds, so the problem keeps its shape from round to round.
        const auto &modules = settings.modules;
        std::vector<std::array<double, 3>> move(modules.size());
        for (size_t m = 0; m < modules.size(); ++m)
            move[m] = {modules[m].dx, modules[m].dy, modules[m].rot_rad};
        // A module counts the spots on it that its frame's lattice indexes where the module starts - not
        // every spot found on it: a frame's other spots (noise, a second lattice, ice) would otherwise make a
        // module the frames barely reach look well determined.
        std::vector<int> module_spots(modules.size(), 0);
        for (const auto &f : frames) {
            const Coord v0 = f.lattice.Vec0(), v1 = f.lattice.Vec1(), v2 = f.lattice.Vec2();
            for (const auto &s : f.spots)
                if (s.module >= 0 && static_cast<size_t>(s.module) < modules.size()) {
                    double x, y;
                    PlaceOnModule(modules[s.module], move[s.module].data(), s.x, s.y, x, y);
                    const Coord recip = nominal_geom.DetectorToRecip(x, y);
                    const double hf = recip * v0, kf = recip * v1, lf = recip * v2;
                    if (std::fabs(hf - std::round(hf)) < 0.15 && std::fabs(kf - std::round(kf)) < 0.15
                        && std::fabs(lf - std::round(lf)) < 0.15)
                        module_spots[s.module]++;
                }
        }
        std::vector<bool> refined(modules.size(), false);
        int refined_count = 0;
        for (size_t m = 0; m < modules.size(); ++m)
            if (settings.refine_modules && module_spots[m] >= settings.min_module_spots) {
                refined[m] = true;
                refined_count++;
            }
        // One module alone has nothing to be placed against: its move is the beam and the orientation.
        if (refined_count < 2)
            std::fill(refined.begin(), refined.end(), false);
        // A module too thin to be placed stays where it was given, whatever its start: a search on a few dozen
        // spots finds noise.
        if (settings.refine_modules)
            for (size_t m = 0; m < modules.size(); ++m)
                if (!refined[m])
                    move[m] = {0, 0, 0};

        // Where nothing but the crystal orientations moves, each spot's residual is reduced to them
        // (OrientationResidual).
        const bool orientation_only = !settings.refine_beam_distance
                                      && std::none_of(refined.begin(), refined.end(), [](bool r) { return r; });

        auto place = [&](const GeomRefineSpot &s, double &x, double &y) {
            if (s.module >= 0 && static_cast<size_t>(s.module) < modules.size())
                PlaceOnModule(modules[s.module], move[s.module].data(), s.x, s.y, x, y);
            else {
                x = s.x;
                y = s.y;
            }
        };

        // Tolerance schedule (fractional-hkl norm), tightening each round like the per-image optimizer.
        auto round_tol = [&](int r) {
            if (settings.rounds <= 1) return 0.2;
            const double t = static_cast<double>(r) / (settings.rounds - 1);
            return 0.3 - (0.3 - 0.1) * t;
        };

        ceres::CauchyLoss loss(loss_scale);

        // Whether the last solve that ran produced a usable solution; folded into result.ok below,
        // which otherwise reports success on spot and frame counts alone.
        bool solve_usable = false;

        for (int round = 0; round < settings.rounds; ++round) {
            EffectiveCell(system, cell_len, cell_ang, eff_len, eff_alpha, eff_beta, eff_gamma);

            double poni_x, poni_y;
            poni(poni_x, poni_y);
            DiffractionGeometry cur = nominal_geom;
            cur.BeamX_pxl(poni_x).BeamY_pxl(poni_y).DetectorDistance_mm(distance_mm)
               .PoniRot1_rad(static_cast<float>(detector_rot[0])).PoniRot2_rad(static_cast<float>(detector_rot[1]));

            const double tol = round_tol(round);
            const double tol_sq = tol * tol;

            // One robust loss for every spot, kept here rather than one allocated (and freed) per spot.
            ceres::Problem::Options problem_options;
            problem_options.loss_function_ownership = ceres::DO_NOT_TAKE_OWNERSHIP;
            ceres::Problem problem(problem_options);
            auto ordering = std::make_shared<ceres::ParameterBlockOrdering>();

            int spots_used = 0, frames_used = 0;
            for (size_t i = 0; i < frames.size(); ++i) {
                const CrystalLattice latt = AngleAxisAndCellToLattice(
                    orient[i].data(), eff_len, eff_alpha, eff_beta, eff_gamma);
                const Coord v0 = latt.Vec0(), v1 = latt.Vec1(), v2 = latt.Vec2();

                int used_in_frame = 0;
                for (const auto &s : frames[i].spots) {
                    double x, y;
                    place(s, x, y);
                    const Coord recip = cur.DetectorToRecip(x, y);
                    const double hf = recip * v0, kf = recip * v1, lf = recip * v2;
                    const double h = std::round(hf), k = std::round(kf), l = std::round(lf);
                    const double dn = (hf - h) * (hf - h) + (kf - k) * (kf - k) + (lf - l) * (lf - l);
                    if (dn > tol_sq)
                        continue;

                    const XtalResidual xr(x, y, lambda, pixel, cos_rot3, sin_rot3, 0.0, h, k, l, system, orientation);
                    if (orientation_only) {
                        auto *r = new OrientationResidual;
                        double b[2];
                        beam_block.Poni(beam, distance_mm, detector_rot, b);
                        xr.ObservedRecip(b, distance_mm, detector_rot, rot_vec, r->recip_obs);
                        xr.PredictedRecipUnrot(cell_len, cell_ang, r->recip_unrot);
                        problem.AddResidualBlock(new ceres::AutoDiffCostFunction<OrientationResidual, 3, 3>(r),
                                                 &loss, orient[i].data());
                    } else if (s.module >= 0 && static_cast<size_t>(s.module) < modules.size() && refined[s.module]) {
                        // The spot as found: the module's move is a block of the residual.
                        const auto &m = modules[s.module];
                        const XtalResidual found(s.x, s.y, lambda, pixel, cos_rot3, sin_rot3, 0.0, h, k, l, system,
                                                 orientation);
                        auto *cost = new SpotCost(found, beam_block, detector_rot, rot_vec, cell_ang, refine_tilt, true,
                                                  s.x - m.x0, s.y - m.y0);
                        if (refine_tilt)
                            problem.AddResidualBlock(cost, &loss, beam_distance, detector_rot, orient[i].data(),
                                                     cell_len, move[s.module].data());
                        else
                            problem.AddResidualBlock(cost, &loss, beam_distance, orient[i].data(), cell_len,
                                                     move[s.module].data());
                    } else {
                        auto *cost = new SpotCost(xr, beam_block, detector_rot, rot_vec, cell_ang, refine_tilt, false);
                        if (refine_tilt)
                            problem.AddResidualBlock(cost, &loss, beam_distance, detector_rot, orient[i].data(),
                                                     cell_len);
                        else
                            problem.AddResidualBlock(cost, &loss, beam_distance, orient[i].data(), cell_len);
                    }
                    ++used_in_frame;
                    ++spots_used;
                }
                if (used_in_frame > 0) {
                    ordering->AddElementToGroup(orient[i].data(), 0);   // eliminate orientations first
                    ++frames_used;
                }
            }

            if (spots_used < settings.min_spots_total) {
                if (round == 0)
                    return result;   // ok stays false
                break;               // keep the previous round's solution
            }

            // Cell-length regularizer (weight ~ coeff * sqrt(N_obs) / L0^2 -> effectively pins the cell).
            const double sqrtN = std::sqrt(static_cast<double>(spots_used));
            if (settings.refine_beam_distance) {
                CellLengthRegularizer *reg = new CellLengthRegularizer{
                    settings.cell_reg_coeff * sqrtN / (cell_len0[0] * cell_len0[0]),
                    settings.cell_reg_coeff * sqrtN / (cell_len0[1] * cell_len0[1]),
                    settings.cell_reg_coeff * sqrtN / (cell_len0[2] * cell_len0[2]),
                    cell_len0[0], cell_len0[1], cell_len0[2]};
                problem.AddResidualBlock(
                    new ceres::AutoDiffCostFunction<CellLengthRegularizer, 3, 3>(reg), nullptr, cell_len);
            }

            // Shared blocks go in the second elimination group (those the problem has; constant ones are
            // stripped by Ceres).
            for (double *b : {beam_distance, detector_rot, cell_len})
                if (problem.HasParameterBlock(b))
                    ordering->AddElementToGroup(b, 1);

            // The refined modules' moves: bounded to what a misplaced module plausibly is, and held to
            // zero in sum by the gauge (ModuleGauge). A 1 px mean shift costs as much there as moving
            // every spot by 1 px would, ten times over.
            std::vector<size_t> moved;
            for (size_t m = 0; m < modules.size(); ++m)
                if (refined[m] && problem.HasParameterBlock(move[m].data()))
                    moved.push_back(m);
            if (moved.size() >= 2) {
                auto *gauge = new ModuleGauge;
                double lever = 0, total = 0;
                for (size_t m : moved)
                    total += module_spots[m];
                for (size_t m : moved) {
                    const double ox = modules[m].x0 - nominal_geom.GetBeamX_pxl();
                    const double oy = modules[m].y0 - nominal_geom.GetBeamY_pxl();
                    const double r = std::max(1.0, std::hypot(ox, oy));
                    gauge->outward.push_back({ox / r, oy / r});
                    gauge->share.push_back(module_spots[m] / total);
                    lever += r;
                }
                gauge->lever_px = lever / static_cast<double>(moved.size());
                gauge->weight = 10.0 * recip_per_px * sqrtN;
                auto *cost = new ceres::DynamicAutoDiffCostFunction<ModuleGauge, 4>(gauge);
                std::vector<double *> blocks;
                for (size_t m : moved) {
                    cost->AddParameterBlock(3);
                    blocks.push_back(move[m].data());
                }
                cost->SetNumResiduals(4);
                problem.AddResidualBlock(cost, nullptr, blocks);
            }
            for (size_t m : moved) {
                ordering->AddElementToGroup(move[m].data(), 1);
                const double start[3] = {modules[m].dx, modules[m].dy, modules[m].rot_rad};
                const double range[3] = {12.0, 12.0, 1.5 * PI / 180.0};
                for (int j = 0; j < 3; ++j) {
                    problem.SetParameterLowerBound(move[m].data(), j, start[j] - range[j]);
                    problem.SetParameterUpperBound(move[m].data(), j, start[j] + range[j]);
                }
            }

            // Beam + distance bounds around the nominal value; cell lengths loosely bounded (the
            // regularizer does the anchoring).
            if (!orientation_only) {
                problem.SetParameterLowerBound(beam_distance, 0, beam0[0] - settings.beam_range_px);
                problem.SetParameterUpperBound(beam_distance, 0, beam0[0] + settings.beam_range_px);
                problem.SetParameterLowerBound(beam_distance, 1, beam0[1] - settings.beam_range_px);
                problem.SetParameterUpperBound(beam_distance, 1, beam0[1] + settings.beam_range_px);
                problem.SetParameterLowerBound(beam_distance, 2,
                    std::max(1.0, nominal_geom.GetDetectorDistance_mm() - settings.distance_range_mm));
                problem.SetParameterUpperBound(beam_distance, 2,
                    nominal_geom.GetDetectorDistance_mm() + settings.distance_range_mm);
                for (int j = 0; j < 3; ++j) {
                    problem.SetParameterLowerBound(cell_len, j, 0.7 * cell_len0[j]);
                    problem.SetParameterUpperBound(cell_len, j, 1.3 * cell_len0[j]);
                }

                if (!settings.refine_beam_distance) {
                    problem.SetParameterBlockConstant(beam_distance);
                    problem.SetParameterBlockConstant(cell_len);
                } else if (!settings.refine_distance)
                    problem.SetManifold(beam_distance, new ceres::SubsetManifold(3, {2}));
            }
            if (refine_tilt) {
                const double tilt_range = 2.0 * PI / 180.0;
                for (int j = 0; j < 2; ++j) {
                    problem.SetParameterLowerBound(detector_rot, j, (j == 0 ? nominal_geom.GetPoniRot1_rad()
                                                                            : nominal_geom.GetPoniRot2_rad()) - tilt_range);
                    problem.SetParameterUpperBound(detector_rot, j, (j == 0 ? nominal_geom.GetPoniRot1_rad()
                                                                            : nominal_geom.GetPoniRot2_rad()) + tilt_range);
                }
            }

            ceres::Solver::Options options;
            options.linear_solver_type = ceres::DENSE_SCHUR;
            options.linear_solver_ordering = ordering;
            options.num_threads = std::max(1, settings.num_threads);
            options.max_num_iterations = 100;
            options.max_solver_time_in_seconds = 60.0;
            options.logging_type = ceres::LoggingType::SILENT;
            options.minimizer_progress_to_stdout = false;

            ceres::Solver::Summary summary;
            ceres::Solve(options, &problem, &summary);
            solve_usable = summary.IsSolutionUsable();

            result.frames_used = frames_used;
            result.spots_used = spots_used;
        }

        // Final fit quality: median detector residual of the retained spots at the refined geometry.
        EffectiveCell(system, cell_len, cell_ang, eff_len, eff_alpha, eff_beta, eff_gamma);
        double poni_x, poni_y;
        poni(poni_x, poni_y);
        DiffractionGeometry cur = nominal_geom;
        cur.BeamX_pxl(poni_x).BeamY_pxl(poni_y).DetectorDistance_mm(distance_mm)
           .PoniRot1_rad(static_cast<float>(detector_rot[0])).PoniRot2_rad(static_cast<float>(detector_rot[1]));
        std::vector<double> residuals;
        constexpr double CAP_PX = 2.0;
        double capped_sum = 0;
        size_t capped_n = 0;
        std::vector<double> module_sq(modules.size(), 0), module_d(modules.size(), 0);
        std::vector<size_t> module_n(modules.size(), 0);
        result.predicted_px.resize(frames.size());
        for (size_t i = 0; i < frames.size(); ++i) {
            const CrystalLattice latt = AngleAxisAndCellToLattice(
                orient[i].data(), eff_len, eff_alpha, eff_beta, eff_gamma);
            const Coord a = latt.Astar(), b = latt.Bstar(), c = latt.Cstar();
            const Coord v0 = latt.Vec0(), v1 = latt.Vec1(), v2 = latt.Vec2();
            for (const auto &s : frames[i].spots) {
                double x, y;
                place(s, x, y);
                const Coord recip = cur.DetectorToRecip(x, y);
                const double hf = recip * v0, kf = recip * v1, lf = recip * v2;
                const double h = std::round(hf), k = std::round(kf), l = std::round(lf);
                const double dn = (hf - h) * (hf - h) + (kf - k) * (kf - k) + (lf - l) * (lf - l);
                const Coord recip_pred = a * (float)h + b * (float)k + c * (float)l;
                const auto [px, py] = cur.RecipToDetector(recip_pred);
                const double res = (std::isfinite(px) && std::isfinite(py)) ? std::hypot(px - x, py - y) : CAP_PX;
                capped_sum += std::min(res, CAP_PX) * std::min(res, CAP_PX);
                capped_n++;
                // The prediction where the spot would be found: the spot's own module's move taken back.
                float fx = NAN, fy = NAN;
                if (std::isfinite(px) && std::isfinite(py)) {
                    fx = static_cast<float>(s.x + (px - x));
                    fy = static_cast<float>(s.y + (py - y));
                }
                result.predicted_px[i].push_back({fx, fy});
                if (s.module >= 0 && static_cast<size_t>(s.module) < modules.size()) {
                    module_sq[s.module] += std::min(res, CAP_PX) * std::min(res, CAP_PX);
                    module_d[s.module] += 1.0 / std::max(1e-6f, recip.Length());
                    module_n[s.module]++;
                }
                if (dn > 0.1 * 0.1)
                    continue;
                if (std::isfinite(px) && std::isfinite(py))
                    residuals.push_back(res);
            }
        }
        if (capped_n > 0)
            result.capped_rms_px = std::sqrt(capped_sum / static_cast<double>(capped_n));
        if (!residuals.empty()) {
            std::nth_element(residuals.begin(), residuals.begin() + residuals.size() / 2, residuals.end());
            result.median_residual_px = residuals[residuals.size() / 2];
        }

        result.ok = solve_usable && result.spots_used >= settings.min_spots_total && result.frames_used > 0;
        result.beam_x_px = poni_x;
        result.beam_y_px = poni_y;
        result.distance_mm = distance_mm;
        result.cell = BuildUnitCell(system, cell_len, cell_ang);
        result.rot1_rad = detector_rot[0];
        result.rot2_rad = detector_rot[1];
        result.module_capped_rms_px.resize(modules.size(), 0);
        result.module_mean_d_A.resize(modules.size(), 0);
        for (size_t m = 0; m < modules.size(); ++m)
            if (module_n[m] > 0) {
                result.module_capped_rms_px[m] = std::sqrt(module_sq[m] / static_cast<double>(module_n[m]));
                result.module_mean_d_A[m] = module_d[m] / static_cast<double>(module_n[m]);
            }
        result.modules = modules;
        for (size_t m = 0; m < modules.size(); ++m) {
            result.modules[m].dx = move[m][0];
            result.modules[m].dy = move[m][1];
            result.modules[m].rot_rad = move[m][2];
            result.modules[m].spots = module_spots[m];
        }
        return result;
    } catch (...) {
        result.ok = false;
        return result;
    }
}

std::vector<ModuleScan> ScanModuleShifts(const DiffractionGeometry &geom,
                                         const std::vector<GeomRefineFrame> &frames,
                                         const std::vector<ModuleShift> &modules,
                                         double range_px, size_t nthreads) {
    // A spot is on its lattice when every fractional index is within this of a whole number.
    constexpr double ON_LATTICE = 0.15;

    struct Spot { float x, y; int frame; };
    std::vector<std::vector<Spot>> on_module(modules.size());
    for (size_t i = 0; i < frames.size(); i++)
        for (const auto *list : {&frames[i].spots, &frames[i].unindexed})
            for (const auto &s : *list)
                if (s.module >= 0 && static_cast<size_t>(s.module) < modules.size())
                    on_module[s.module].push_back({s.x, s.y, static_cast<int>(i)});

    std::vector<std::array<Coord, 3>> axes(frames.size());
    for (size_t i = 0; i < frames.size(); i++)
        axes[i] = {frames[i].lattice.Vec0(), frames[i].lattice.Vec1(), frames[i].lattice.Vec2()};

    auto score = [&](const std::vector<Spot> &spots, const ModuleShift &m, double dx, double dy) {
        int n = 0;
        for (const auto &s : spots) {
            double x, y;
            const double move[3] = {m.dx + dx, m.dy + dy, m.rot_rad};
            PlaceOnModule(m, move, s.x, s.y, x, y);
            const Coord recip = geom.DetectorToRecip(x, y);
            bool on = true;
            for (const auto &v : axes[s.frame]) {
                const double f = recip * v;
                on &= std::fabs(f - std::round(f)) < ON_LATTICE;
            }
            n += on;
        }
        return n;
    };

    std::vector<ModuleScan> out(modules.size());
    for (size_t m = 0; m < modules.size(); m++) {
        auto &r = out[m];
        r.on_lattice_at_zero = r.on_lattice_at_best = score(on_module[m], modules[m], 0, 0);
        // Coarse: whole-pixel steps over the range; then half-pixel steps around the best of them.
        for (int pass = 0; pass < 2; pass++) {
            const double step = pass == 0 ? 1.0 : 0.5;
            const double half = pass == 0 ? range_px : 1.0;
            const double cx = r.dx, cy = r.dy;
            std::vector<double> offsets;
            for (double d = -half; d <= half; d += step)
                offsets.push_back(d);
            // Every shift of the grid is scored on the workers, then the grid is walked in order: the first
            // shift with the most spots wins, as it would on one thread.
            const int n = static_cast<int>(offsets.size());
            std::vector<int> on_lattice(n * n);
            ParallelFor(n * n, nthreads, [&](int k) {
                on_lattice[k] = score(on_module[m], modules[m], cx + offsets[k % n], cy + offsets[k / n]);
            });
            for (int k = 0; k < n * n; k++)
                if (on_lattice[k] > r.on_lattice_at_best) {
                    r.on_lattice_at_best = on_lattice[k];
                    r.dx = cx + offsets[k % n];
                    r.dy = cy + offsets[k / n];
                }
        }
    }
    return out;
}