// 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 "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];
}

// XtalResidual for a spot on a module that is being placed: block `move` = (dx, dy, rot) of the module,
// (ux, uy) the spot's offset from the module centre. Moving the spot by R(rot) u - u + (dx, dy) 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.
struct ModuleXtalResidual {
    XtalResidual residual;
    double ux, uy;

    template <typename T>
    bool operator()(const T *const beam, const T *const distance_mm, const T *const detector_rot,
                    const T *const rotation_axis, const T *const p0, const T *const p1, const T *const p2,
                    const T *const move, T *out) const {
        using std::cos;
        using std::sin;
        const T c = cos(move[2]), s = sin(move[2]);
        const T moved_beam[2] = {beam[0] - ((c - T(1)) * T(ux) - s * T(uy) + move[0]),
                                 beam[1] - (s * T(ux) + (c - T(1)) * T(uy) + move[1])};
        return residual(moved_beam, distance_mm, detector_rot, rotation_axis, p0, p1, p2, out);
    }
};

// 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.
//
// With the tilt refined, the two tilt-like patterns are held at zero as well: a detector tilted by a small
// angle about y moves a point (x, y) (from the beam) by (x^2, x y) / D times the angle, one about x by
// (x y, y^2) / D. Each is taken as the modules' moves projected on that field, in units of the lever.
struct ModuleGauge {
    std::vector<std::array<double, 2>> outward;   // unit vector from the beam to each module centre
    std::vector<std::array<double, 2>> centre;    // each module centre from the beam, over lever_px
    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;
    bool tilt = false;

    template <typename T>
    bool operator()(T const *const *move, T *out) const {
        T sx(0), sy(0), sr(0), so(0), t1(0), t2(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]));
            const double x = centre[i][0], y = centre[i][1];
            t1 += w * (move[i][0] * T(x * x) + move[i][1] * T(x * y));
            t2 += w * (move[i][0] * T(x * y) + move[i][1] * T(y * y));
        }
        out[0] = T(weight) * sx;
        out[1] = T(weight) * sy;
        out[2] = T(weight) * sr;
        out[3] = T(weight) * so;
        out[4] = tilt ? T(weight) * t1 : T(0);
        out[5] = tilt ? T(weight) * t2 : T(0);
        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). Detector tilt and the rotation axis are constant.
        double beam[2] = {nominal_geom.GetBeamX_pxl(), nominal_geom.GetBeamY_pxl()};
        double distance_mm = nominal_geom.GetDetectorDistance_mm();
        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();

        // 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};

        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;
        };

        // 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);

            DiffractionGeometry cur = nominal_geom;
            cur.BeamX_pxl(beam[0]).BeamY_pxl(beam[1]).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;

            ceres::Problem problem;
            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;

                    if (s.module >= 0 && static_cast<size_t>(s.module) < modules.size() && refined[s.module]) {
                        const auto &m = modules[s.module];
                        problem.AddResidualBlock(
                            new ceres::AutoDiffCostFunction<ModuleXtalResidual, 3, 2, 1, 2, 3, 3, 3, 3, 3>(
                                new ModuleXtalResidual{XtalResidual(s.x, s.y, lambda, pixel, cos_rot3, sin_rot3,
                                                                    0.0, h, k, l, system, orientation),
                                                       s.x - m.x0, s.y - m.y0}),
                            new ceres::CauchyLoss(loss_scale),
                            beam, &distance_mm, detector_rot, rot_vec, orient[i].data(), cell_len, cell_ang,
                            move[s.module].data());
                    } else
                        problem.AddResidualBlock(
                            new ceres::AutoDiffCostFunction<XtalResidual, 3, 2, 1, 2, 3, 3, 3, 3>(
                                new XtalResidual(x, y, lambda, pixel, cos_rot3, sin_rot3, 0.0,
                                                 h, k, l, system, orientation)),
                            new ceres::CauchyLoss(loss_scale),
                            beam, &distance_mm, detector_rot, rot_vec, orient[i].data(), cell_len, cell_ang);
                    ++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; constant ones are stripped by Ceres.
            for (double *b : {beam, &distance_mm, detector_rot, rot_vec, cell_len, cell_ang})
                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->centre.push_back({ox, oy});
                    gauge->share.push_back(module_spots[m] / total);
                    lever += r;
                }
                gauge->lever_px = lever / static_cast<double>(moved.size());
                for (auto &c : gauge->centre)
                    c = {c[0] / gauge->lever_px, c[1] / gauge->lever_px};
                gauge->weight = 10.0 * recip_per_px * sqrtN;
                gauge->tilt = settings.refine_tilt;
                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(6);
                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).
            problem.SetParameterLowerBound(beam, 0, nominal_geom.GetBeamX_pxl() - settings.beam_range_px);
            problem.SetParameterUpperBound(beam, 0, nominal_geom.GetBeamX_pxl() + settings.beam_range_px);
            problem.SetParameterLowerBound(beam, 1, nominal_geom.GetBeamY_pxl() - settings.beam_range_px);
            problem.SetParameterUpperBound(beam, 1, nominal_geom.GetBeamY_pxl() + settings.beam_range_px);
            problem.SetParameterLowerBound(&distance_mm, 0,
                std::max(1.0, nominal_geom.GetDetectorDistance_mm() - settings.distance_range_mm));
            problem.SetParameterUpperBound(&distance_mm, 0,
                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_distance)
                problem.SetParameterBlockConstant(&distance_mm);
            if (settings.refine_tilt && settings.refine_beam_distance) {
                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);
                }
            } else
                problem.SetParameterBlockConstant(detector_rot);
            problem.SetParameterBlockConstant(rot_vec);
            problem.SetParameterBlockConstant(cell_ang);       // angles not refined in v1
            if (!settings.refine_beam_distance) {
                problem.SetParameterBlockConstant(beam);
                problem.SetParameterBlockConstant(&distance_mm);
                problem.SetParameterBlockConstant(cell_len);
            }

            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);
        DiffractionGeometry cur = nominal_geom;
        cur.BeamX_pxl(beam[0]).BeamY_pxl(beam[1]).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 = beam[0];
        result.beam_y_px = beam[1];
        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) {
    // 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;
            for (double dy = -half; dy <= half; dy += step)
                for (double dx = -half; dx <= half; dx += step) {
                    const int n = score(on_module[m], modules[m], cx + dx, cy + dy);
                    if (n > r.on_lattice_at_best) {
                        r.on_lattice_at_best = n;
                        r.dx = cx + dx;
                        r.dy = cy + dy;
                    }
                }
        }
    }
    return out;
}