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

#include "RigidBodyRefine.h"

#include <algorithm>
#include <array>
#include <chrono>
#include <cmath>
#include <complex>
#include <future>
#include <memory>
#include <vector>

#include <Eigen/Dense>
#include <ceres/ceres.h>
#include <ceres/rotation.h>

#include "gemmi/dencalc.hpp"   // DensityCalculator
#include "gemmi/it92.hpp"      // IT92 x-ray form factors
#include "gemmi/scaling.hpp"   // Scaling (bulk solvent + anisotropic B)

#include "ModelFFT.h"               // MapToFPhi
#include "ModelGrid.h"              // PutModelDensityOnGrid, PutMaskOnGrid
#include "ModelScaling.h"           // FitModelScale
#ifdef JFJOCH_USE_CUDA
#include "RigidBodyGPU.h"           // RigidBodyTargetGPU
#endif
#include "../../common/JFJochMath.h"   // PI
#include "../../common/ParallelFor.h"
#include "../../common/Logger.h"

namespace {

using Table = gemmi::IT92<float>;

// The ladder the placement is walked down. It starts coarse because the model arrives already placed
// but out by a cell's worth of non-isomorphism: at 6 A a few hundred reflections see the body as a
// blob and the target has one broad minimum, and each finer zone starts from the previous one's
// answer. It stops at 3.5 A, which is where rigid-body refinement is conventionally run (it is
// REFMAC's own default through dimple) - the movement being recovered is a few tenths of an
// angstrom, a tenth of that resolution, so it is well determined there, while a finer zone costs
// (1/d)^3 in grid points and reflections for a placement it cannot meaningfully sharpen.
constexpr double LADDER[] = {6.0, 4.5, 3.5};

// The step of the forward-difference rotation columns of the Jacobian, as a fraction of the zone's
// resolution - so it is 0.06 A of atom displacement at 6 A and 0.035 A at 3.5 A. A step fixed in
// angstroms instead is far too small for the coarse zones, where a structure factor barely notices it
// and the derivative is swallowed by rounding: measured, a fixed 0.02 A left the 6 A zone at 0.35
// degrees where this rule takes it to 2.79, which is most of the way to the answer.
constexpr double JACOBIAN_STEP_FRACTION = 0.01;

// How finely a zone's maps are sampled: DensityCalculator's rate, for a spacing of d_min / (2 * rate).
constexpr double GRID_RATE = 1.5;

// Ceres' own numeric differentiation steps by |x| * relative_step_size, which is zero at the start of
// every zone (the placement begins at no shift), so the Jacobian is supplied by RigidBodyTarget.
class RigidBodyCost : public ceres::CostFunction {
public:
    explicit RigidBodyCost(RigidBodyTargetBase &target) : target_(target) {
        set_num_residuals(static_cast<int>(target.NumObservations()));
        mutable_parameter_block_sizes()->push_back(6);
    }

    bool Evaluate(double const *const *parameters, double *residuals, double **jacobians) const override {
        // Ceres evaluates every step it accepts twice: the residuals alone to decide whether to take
        // it, then residuals and Jacobian together at the same point. The residuals are a function of
        // the placement alone, so the second time they are the first time's, copied.
        if (!last_residuals_.empty() && std::equal(parameters[0], parameters[0] + 6, last_q_.begin())) {
            std::copy(last_residuals_.begin(), last_residuals_.end(), residuals);
        } else {
            if (!target_.Residuals(parameters[0], residuals))
                return false;
            std::copy(parameters[0], parameters[0] + 6, last_q_.begin());
            last_residuals_.assign(residuals, residuals + num_residuals());
        }
        if (jacobians != nullptr && jacobians[0] != nullptr)
            return target_.Jacobian(parameters[0], jacobians[0]);
        return true;
    }

private:
    RigidBodyTargetBase &target_;
    mutable std::array<double, 6> last_q_{};
    mutable std::vector<double> last_residuals_;   // at last_q_; empty until an evaluation succeeds
};

// The directions in which this space group's origin is free. Translating the whole cell content along
// one of them multiplies every F by a phase and leaves every |F| EXACTLY unchanged, so the target
// cannot determine that component: all three directions in P1, the unique axis in a polar group. The
// R-free gate cannot stand in for this - it is a function of |F| too, so along such a direction it
// sees only grid noise and commits or not by coin flip, while the other five parameters carry the
// noise in with them. The free directions are the common fixed subspace of the group's rotation
// parts, and the projector onto it is simply their average.
}  // namespace

gemmi::Mat33 RigidBodyGaugeProjector(const gemmi::SpaceGroup &sg, const gemmi::UnitCell &cell) {
    const gemmi::GroupOps gops = sg.operations();
    double m[3][3] = {};
    const double n = static_cast<double>(gops.sym_ops.size()) * gemmi::Op::DEN;
    for (const gemmi::Op &op : gops.sym_ops)
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                m[i][j] += static_cast<double>(op.rot[i][j]) / n;
    const gemmi::Mat33 mean(m[0][0], m[0][1], m[0][2],
                            m[1][0], m[1][1], m[1][2],
                            m[2][0], m[2][1], m[2][2]);
    // Fractional projector taken into orthogonal space, where the parameters live.
    return cell.orth.mat.multiply(mean).multiply(cell.frac.mat);
}

// An empty grid of the zone's size, the size DensityCalculator gives its own at GRID_RATE.
gemmi::Grid<float> RigidBodyZoneGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) {
    gemmi::Grid<float> grid;
    grid.unit_cell = cell;
    grid.spacegroup = &sg;
    grid.set_size_from_spacing(d_min / (2 * GRID_RATE), gemmi::GridSizeRounding::Up);
    return grid;
}

double RigidBodyJacobianStep(double d_min) {
    return JACOBIAN_STEP_FRACTION * d_min;
}

std::vector<double> RigidBodyLadder(double d_min) {
    std::vector<double> ladder;
    for (double zone : LADDER)
        if (zone >= d_min)
            ladder.push_back(zone);
    if (ladder.empty())
        ladder.push_back(d_min);
    return ladder;
}

std::vector<gemmi::Position> ModelPositions(const gemmi::Model &model) {
    std::vector<gemmi::Position> pos;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            for (const gemmi::Atom &a : r.atoms)
                pos.push_back(a.pos);
    return pos;
}

void SetModelPositions(gemmi::Model &model, const std::vector<gemmi::Position> &pos) {
    size_t i = 0;
    for (gemmi::Chain &ch : model.chains)
        for (gemmi::Residue &r : ch.residues)
            for (gemmi::Atom &a : r.atoms)
                a.pos = pos[i++];
}

SymmetryComposition::SymmetryComposition(const gemmi::Grid<float> &grid, double d_min,
                                         const std::vector<gemmi::Miller> &hkl) {
    const gemmi::UnitCell &cell = grid.unit_cell;
    const gemmi::GroupOps gops = grid.spacegroup->operations();
    const gemmi::ReciprocalAsu asu(grid.spacegroup);
    // prepare_asu_data()'s box on the half-l transform of this grid, (nu, nv, nw/2 + 1): up to the
    // Nyquist frequency but not on it along u and v, and to nw/2 along w.
    const gemmi::Miller lim = cell.get_hkl_limits(d_min);
    const int max_h = std::min((grid.nu - 1) / 2, lim[0]);
    const int max_k = std::min((grid.nv - 1) / 2, lim[1]);
    const int max_l = std::min(grid.nw / 2, lim[2]);
    const double max_1_d2 = 1.0 / (d_min * d_min);
    ops_ = gops.sym_ops.size();
    centring_ = static_cast<double>(gops.cen_ops.size());
    row_.assign(hkl.size(), -1);
    for (size_t i = 0; i < hkl.size(); i++) {
        const gemmi::Miller &h = hkl[i];
        const double inv_d2 = cell.calculate_1_d2(h);
        if (std::abs(h[0]) > max_h || std::abs(h[1]) > max_k || std::abs(h[2]) > max_l ||
            !(inv_d2 < max_1_d2) || !asu.is_in(h) || gops.is_systematically_absent(h) ||
            (h[0] == 0 && h[1] == 0 && h[2] == 0))
            continue;
        row_[i] = static_cast<int>(hkl_.size());
        hkl_.push_back(h);
        inv_d2_.push_back(inv_d2);
        for (const gemmi::Op &op : gops.sym_ops) {
            const gemmi::Miller k = op.apply_to_hkl(h);
            Term t;
            t.conj = k[2] < 0;   // the half-l grid holds l >= 0, and F1(-k) = conj F1(k) for a real map
            const int ku = t.conj ? -k[0] : k[0], kv = t.conj ? -k[1] : k[1], kw = t.conj ? -k[2] : k[2];
            t.index = gemmi::modulo(ku, grid.nu) +
                      static_cast<size_t>(grid.nu) * (gemmi::modulo(kv, grid.nv) + static_cast<size_t>(grid.nv) * kw);
            t.phase = std::polar(1.0, -op.phase_shift(h));   // gemmi's phase_shift is -2 pi h.t
            t.s = cell.frac.mat.left_multiply(gemmi::Vec3(k[0], k[1], k[2]));
            t.k = k;
            terms_.push_back(t);
        }
    }
}

void SymmetryComposition::Compose(const gemmi::FPhiGrid<float> &f1, double unblur,
                                  std::vector<std::complex<double>> &f,
                                  std::vector<std::array<std::complex<double>, 3>> *df_dt,
                                  size_t nthreads) const {
    const int n = static_cast<int>(hkl_.size());
    f.assign(n, 0.0);
    if (df_dt != nullptr)
        df_dt->assign(n, {});
    const std::complex<double> two_pi_i(0.0, 2 * PI);
    ParallelChunks(n, nthreads, [&](int lo, int hi) {
        for (int m = lo; m < hi; m++) {
            std::complex<double> sum = 0;
            std::array<std::complex<double>, 3> d{};
            for (size_t o = 0; o < ops_; o++) {
                const Term &t = terms_[m * ops_ + o];
                const std::complex<double> value = f1.data[t.index];
                const std::complex<double> term = t.phase * (t.conj ? std::conj(value) : value);
                sum += term;
                d[0] += term * t.s.x;
                d[1] += term * t.s.y;
                d[2] += term * t.s.z;
            }
            // prepare_asu_data()'s unblur: exp(B_blur |s|^2 / 4)
            const double scale = centring_ * std::exp(unblur * 0.25 * inv_d2_[m]);
            f[m] = scale * sum;
            if (df_dt != nullptr)
                for (int k = 0; k < 3; k++)
                    (*df_dt)[m][k] = scale * two_pi_i * d[k];
        }
    });
}

RigidBodyTargetBase::RigidBodyTargetBase(const gemmi::Model &model) : base_(ModelPositions(model)) {
    if (base_.empty())
        return;
    for (const gemmi::Position &p : base_)
        centre_ += p;
    centre_ *= 1.0 / static_cast<double>(base_.size());
    double r2 = 0;
    for (const gemmi::Position &p : base_)
        r2 += centre_.dist_sq(p);
    rms_radius_ = std::sqrt(r2 / static_cast<double>(base_.size()));
}

RigidBodyTarget::RigidBodyTarget(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                 size_t nthreads)
    : RigidBodyTargetBase(model), model_(model), cell_(cell), sg_(sg), nthreads_(nthreads),
      column_models_(3, model) {}

// Parameters are carried as six lengths in angstroms - the first three are the angle-axis rotation
// vector multiplied by the model's rms radius, so a unit of each of the six moves a typical atom by
// the same amount. The rotation is about the model centroid, which decorrelates it from the
// translation, and is applied to the positions only: an anisotropic U is not turned with the body.
void RigidBodyTargetBase::Place(const double q[6], gemmi::Model &model) const {
    const double aa[3] = {q[0] / rms_radius_, q[1] / rms_radius_, q[2] / rms_radius_};
    size_t i = 0;
    for (gemmi::Chain &ch : model.chains)
        for (gemmi::Residue &r : ch.residues)
            for (gemmi::Atom &a : r.atoms) {
                const double p[3] = {base_[i].x - centre_.x, base_[i].y - centre_.y, base_[i].z - centre_.z};
                double rp[3];
                ceres::AngleAxisRotatePoint(aa, p, rp);
                a.pos = gemmi::Position(rp[0] + centre_.x + q[3],
                                        rp[1] + centre_.y + q[4],
                                        rp[2] + centre_.z + q[5]);
                ++i;
            }
}

void RigidBodyTarget::SetZone(const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs, double d_min) {
    fobs_ = fobs;
    d_min_ = d_min;
    double sum = 0;
    for (const auto &hv : fobs_.v)
        sum += hv.value.value;
    f_mean_ = fobs_.v.empty() ? 1.0 : sum / static_cast<double>(fobs_.v.size());
    solvent_fitted_ = false;
    have_point_ = false;
    const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell_, sg_, d_min);
    orbit_leaders_ = OrbitLeaders(grid, nthreads_);
    std::vector<gemmi::Miller> hkl;
    for (const auto &hv : fobs_.v)
        hkl.push_back(hv.hkl);
    composition_.emplace(grid, d_min, hkl);
}

// The model's Fcalc at the zone's composed indices, from one copy of it on the zone's grid.
void RigidBodyTarget::CopyFcalc(const gemmi::Model &model, std::vector<std::complex<double>> &f,
                                std::vector<std::array<std::complex<double>, 3>> *df_dt) const {
    gemmi::DensityCalculator<Table, float> dc;
    dc.d_min = d_min_;
    dc.rate = GRID_RATE;
    dc.grid.unit_cell = cell_;
    dc.grid.spacegroup = &sg_;
    dc.set_refmac_compatible_blur(model);
    PutModelDensityOnGrid(dc, model, {}, nthreads_);   // no orbits: the copy is not symmetrized
    composition_->Compose(MapToFPhi(dc.grid), dc.blur, f, df_dt, nthreads_);
}

// One target evaluation: place the model, recompute Fcalc and the bulk-solvent mask to the zone's
// resolution, re-fit the scale, and hand back the amplitude residuals.
bool RigidBodyTarget::Residuals(const double q[6], double *residuals) {
    ++evaluations;
    Place(q, model_);
    std::vector<std::complex<double>> fc;
    std::vector<std::array<std::complex<double>, 3>> dfc_dt;
    CopyFcalc(model_, fc, &dfc_dt);

    const std::vector<gemmi::Miller> &hkl = composition_->Hkl();
    std::vector<std::complex<float>> fmask;
    if (hold_mask && have_point_) {
        fmask = fmask_;
    } else {
        // The mask is symmetric as gridded, so it is read at h directly.
        gemmi::Grid<float> mask_grid = RigidBodyZoneGrid(cell_, sg_, d_min_);
        PutMaskOnGrid(mask_grid, model_, orbit_leaders_, nthreads_);
        const gemmi::FPhiGrid<float> fm = MapToFPhi(mask_grid);
        for (const gemmi::Miller &h : hkl)
            fmask.push_back(fm.get_value_by_hkl(h));
    }
    if (hkl.empty())
        return false;
    gemmi::AsuData<std::complex<float>> fcalc, fmask_data;
    for (size_t m = 0; m < hkl.size(); m++) {
        fcalc.v.push_back({hkl[m], std::complex<float>(fc[m])});
        fmask_data.v.push_back({hkl[m], fmask[m]});
    }

    // Re-fitted at every evaluation: with the scale held at the starting placement's value the
    // target would measure the scale as much as the placement, and the body would translate to
    // repair a scale error instead of moving where the density is.
    //
    // The bulk solvent is not part of that scale. k_sol and b_sol describe the disordered solvent
    // of the crystal rather than the fit of one placement, so they are fitted once per zone - by
    // FitModelScale, inside the same physical box the reported fit is searched in - and then held
    // while the overall scale and the anisotropic B follow the body. Leaving them free at every
    // evaluation, which is what gemmi's unbounded Levenberg-Marquardt did here, puts a solvent
    // term with no physical meaning inside the target that decides where the model goes: measured
    // over a corpus of deposited models, 40% of the evaluations came out with b_sol outside
    // 10-80 A^2, some of them negative, which is a solvent that GROWS with resolution.
    gemmi::Scaling<float> scaling(cell_, &sg_);
    scaling.use_solvent = true;
    scaling.prepare_points(fcalc, fobs_, &fmask_data);
    if (scaling.points.empty())
        return false;
    if (!solvent_fitted_) {
        FitModelScale(scaling, {}, nthreads_);
        k_sol = scaling.k_sol;
        b_sol = scaling.b_sol;
        solvent_fitted_ = true;
    }
    scaling.k_sol = k_sol;
    scaling.b_sol = b_sol;
    scaling.fix_k_sol = true;
    scaling.fix_b_sol = true;
    scaling.fit_isotropic_b_approximately();
    scaling.fit_parameters();
    scaling.scale_data(fcalc, &fmask_data);

    // An observation with no calculated amplitude gets residual 0, which drops it from the target
    // rather than scoring it as a perfect fit: its Jacobian row comes out zero as well, and the set
    // that matches is fixed by the cell, the group and the zone, so it does not move as the body
    // does. It is counted and reported because a large count is a statement about the model rather
    // than about this refinement - a group whose reflection conditions the data do not obey leaves
    // half of them with nothing to compare against.
    const std::vector<int> &row = composition_->Row();
    unmatched = 0;
    for (size_t i = 0; i < fobs_.v.size(); ++i) {
        if (row[i] < 0)
            ++unmatched;
        residuals[i] = row[i] < 0 ? 0.0 : (fobs_.v[i].value.value - std::abs(fcalc.v[row[i]].value)) / f_mean_;
    }

    have_point_ = true;
    std::copy(q, q + 6, q_.begin());
    fc_ = std::move(fc);
    dfc_dt_ = std::move(dfc_dt);
    fmask_ = std::move(fmask);
    k_overall_ = scaling.k_overall;
    b_star_ = scaling.b_star;
    return true;
}

// The Jacobian at the last evaluation's placement, without repeating it. The derivative of Fcalc
// with respect to the translation came with that evaluation, exact. The rotation's is a forward
// difference, one step per axis, each needing only one copy of the model gridded and transformed:
// the three run in parallel, column j on its own copy of the model, on threads of their own rather
// than on ParallelFor's pool, for the reason the null's replicates do (ModelValidation.cpp) - a
// gridding spreads over the pool, which it cannot do from a pool worker, where a parallel pass runs
// inline. The bulk-solvent mask is held at the evaluation's: the residuals see it move with the
// body, the Jacobian does not, which makes the steps Levenberg-Marquardt proposes slightly different
// but not the cost it accepts them on.
bool RigidBodyTarget::Jacobian(const double q[6], double *jacobian) {
    if (!have_point_ || !std::equal(q, q + 6, q_.begin())) {
        std::vector<double> residuals(NumObservations());
        if (!Residuals(q, residuals.data()))
            return false;
    }
    ++jacobians;
    const double step = RigidBodyJacobianStep(d_min_);
    std::array<std::vector<std::complex<double>>, 3> shifted;
    const std::launch policy = nthreads_ > 1 ? std::launch::async : std::launch::deferred;
    std::vector<std::future<void>> running;
    for (int j = 0; j < 3; j++)
        running.push_back(std::async(policy, [&, j] {
            double qj[6];
            std::copy(q_.begin(), q_.end(), qj);
            qj[j] += step;
            Place(qj, column_models_[j]);
            CopyFcalc(column_models_[j], shifted[j], nullptr);
        }));
    for (std::future<void> &f : running)
        f.get();

    // d|F_scaled|/dq at the evaluation's scale: F_scaled = K(h) (Fcalc + k_sol exp(-b_sol s^2) Fmask),
    // so the derivative of its amplitude is K Re(conj(F_total) dFcalc/dq) / |F_total|. The scale's
    // own derivatives are gemmi's, for the parameters it re-fits: k_overall and the constrained B*.
    gemmi::Scaling<float> scaling(cell_, &sg_);
    scaling.use_solvent = true;
    scaling.fix_k_sol = true;
    scaling.fix_b_sol = true;
    scaling.k_sol = k_sol;
    scaling.b_sol = b_sol;
    scaling.k_overall = k_overall_;
    scaling.b_star = b_star_;
    const size_t n = NumObservations();
    const size_t p = 1 + scaling.constraint_matrix.size();
    std::vector<double> jk(n * p, 0.0);
    std::vector<double> dy_da(p);
    const std::vector<int> &row = composition_->Row();
    const std::vector<gemmi::Miller> &hkl = composition_->Hkl();
    std::fill(jacobian, jacobian + n * 6, 0.0);
    for (size_t i = 0; i < n; i++) {
        const int m = row[i];
        if (m < 0)
            continue;
        const gemmi::Scaling<float>::Point point{hkl[m], cell_.calculate_stol_sq(hkl[m]),
                                                 std::complex<float>(fc_[m]), fmask_[m],
                                                 fobs_.v[i].value.value, fobs_.v[i].value.sigma};
        scaling.compute_value_and_derivatives(point, dy_da);
        for (size_t a = 0; a < p; a++)
            jk[i * p + a] = -dy_da[a] / f_mean_;
        const std::complex<double> f_total(scaling.get_fcalc(point));
        const double k = scaling.get_overall_scale_factor(hkl[m]) / (std::abs(f_total) * f_mean_);
        for (int j = 0; j < 6; j++) {
            const std::complex<double> dfc = j < 3 ? (shifted[j][m] - fc_[m]) / step : dfc_dt_[m][j - 3];
            jacobian[i * 6 + j] = -k * std::real(std::conj(f_total) * dfc);
        }
    }

    // The residuals are taken at the scale re-fitted to each placement, so their Jacobian is the
    // fixed-scale one with the part a scale change absorbs projected out: J - J_k (J_k^T J_k)^-1 J_k^T J,
    // J_k being the residuals' derivatives with respect to the scale parameters. Without it the steps
    // would be those of a target whose scale stays put, overstating the curvature along any direction
    // a scale change could follow.
    // Following Golub & Pereyra (1973) SIAM J. Numer. Anal. 10, 413-432, in the form of Kaufman (1975) BIT 15, 49-57
    Eigen::MatrixXd jtj = Eigen::MatrixXd::Zero(p, p), jtq = Eigen::MatrixXd::Zero(p, 6);
    for (size_t i = 0; i < n; i++)
        for (size_t a = 0; a < p; a++) {
            for (size_t b = 0; b < p; b++)
                jtj(a, b) += jk[i * p + a] * jk[i * p + b];
            for (int j = 0; j < 6; j++)
                jtq(a, j) += jk[i * p + a] * jacobian[i * 6 + j];
        }
    const Eigen::MatrixXd x = jtj.ldlt().solve(jtq);
    for (size_t i = 0; i < n; i++)
        for (int j = 0; j < 6; j++)
            for (size_t a = 0; a < p; a++)
                jacobian[i * 6 + j] -= jk[i * p + a] * x(a, j);
    return true;
}

// One rigid body, not groups: a fragment-screening model arrives already solved and isomorphous, and
// the movement to recover is the crystal's, not the molecule's. Splitting it into domains or giving a
// bound ligand its own six parameters would refine against evidence this data does not separately
// carry, and the ligand is what the difference map is meant to show rather than model away.
double RigidBodyReachDeg(const gemmi::Model &model, double d_min) {
    const std::vector<gemmi::Position> pos = ModelPositions(model);
    if (pos.empty())
        return 0.0;
    gemmi::Position centre;
    for (const gemmi::Position &p : pos)
        centre += p;
    centre *= 1.0 / static_cast<double>(pos.size());
    double r2 = 0;
    for (const gemmi::Position &p : pos)
        r2 += centre.dist_sq(p);
    const double rms_radius = std::sqrt(r2 / static_cast<double>(pos.size()));
    // The first zone RefineRigidBody walks, which is d_min itself where d_min is coarser than all
    // of the ladder.
    const double first_zone = std::max(LADDER[0], d_min);
    return first_zone / rms_radius * 180.0 / PI;
}

//
// Some of the translation would be a gauge rather than a quantity - the origin is free in all three
// directions in P1 and along the unique axis in a polar group, and |F| does not change when the whole
// content moves along it - so that component is projected out after every zone. Neither of the two
// things that might look like they cover it actually does: the R-free gate is a function of |F| and
// therefore blind to exactly this, and the LM damping follows the gauge column of the Jacobian, which
// is not zero but noise divided by the difference step.
RigidBodyRefineResult RefineRigidBody(gemmi::Model &model,
                                      const gemmi::UnitCell &cell,
                                      const gemmi::SpaceGroup &sg,
                                      const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs,
                                      double d_min,
                                      Logger &logger,
                                      size_t nthreads,
                                      RigidBodyGPUPool *gpu) {
    const auto t0 = std::chrono::steady_clock::now();
    RigidBodyRefineResult result;
    if (fobs.v.empty())
        return result;
    std::unique_ptr<RigidBodyTargetBase> target_backend;
#ifdef JFJOCH_USE_CUDA
    if (gpu != nullptr)
        target_backend = std::make_unique<RigidBodyTargetGPU>(*gpu, model, cell, sg, nthreads);
#endif
    if (!target_backend)
        target_backend = std::make_unique<RigidBodyTarget>(model, cell, sg, nthreads);
    RigidBodyTargetBase &target = *target_backend;
    if (!target.Usable())
        return result;

    const std::vector<double> ladder = RigidBodyLadder(d_min);
    const gemmi::Mat33 gauge = RigidBodyGaugeProjector(sg, cell);
    double q[6] = {0, 0, 0, 0, 0, 0};
    bool any_zone_solved = false;
    for (double zone : ladder) {
        gemmi::AsuData<gemmi::ValueSigma<float>> zone_obs;
        zone_obs.unit_cell_ = fobs.unit_cell_;
        zone_obs.spacegroup_ = fobs.spacegroup_;
        for (const auto &hv : fobs.v)
            if (cell.calculate_d(hv.hkl) >= zone)
                zone_obs.v.push_back(hv);
        if (zone_obs.v.size() < 50)
            continue;
        result.zones.push_back(zone);   // the ladder WALKED, which a thin zone drops out of
        target.SetZone(zone_obs, zone);

        ceres::Problem problem;
        problem.AddResidualBlock(new RigidBodyCost(target), nullptr, q);
        ceres::Solver::Options options;
        options.linear_solver_type = ceres::DENSE_QR;
        options.max_num_iterations = 15;
        options.function_tolerance = 1e-4;
        options.parameter_tolerance = 1e-4;
        options.logging_type = ceres::LoggingType::SILENT;
        ceres::Solver::Summary summary;
        const int evaluations_before = target.evaluations, jacobians_before = target.jacobians;
        const auto zone_t0 = std::chrono::steady_clock::now();
        ceres::Solve(options, &problem, &summary);
        any_zone_solved = any_zone_solved || summary.IsSolutionUsable();
        const gemmi::Vec3 along = gauge.multiply(gemmi::Vec3(q[3], q[4], q[5]));
        q[3] -= along.x; q[4] -= along.y; q[5] -= along.z;
        logger.Debug("Rigid body zone {:.1f} A: {} reflections ({} without a calculated amplitude), "
                     "solvent k_sol {:.2f} b_sol {:.0f} A^2, {} iterations, {} evaluations, {} Jacobians, "
                     "{:.2f} s, rotation {:.3f} deg, translation {:.3f} A", zone, zone_obs.v.size(),
                     target.unmatched, target.k_sol, target.b_sol,
                     summary.iterations.empty() ? 0 : summary.iterations.size() - 1,
                     target.evaluations - evaluations_before, target.jacobians - jacobians_before,
                     std::chrono::duration<double>(std::chrono::steady_clock::now() - zone_t0).count(),
                     std::sqrt(q[0]*q[0] + q[1]*q[1] + q[2]*q[2]) / target.RmsRadius() * 180.0 / PI,
                     std::sqrt(q[3]*q[3] + q[4]*q[4] + q[5]*q[5]));
    }

    target.Place(q, model);   // Ceres left the model at its last trial; put it at the answer
    result.evaluations = target.evaluations;
    result.jacobians = target.jacobians;
    result.converged = any_zone_solved;
    result.k_sol = target.k_sol;
    result.b_sol = target.b_sol;
    const double aa = std::sqrt(q[0] * q[0] + q[1] * q[1] + q[2] * q[2]) / target.RmsRadius();
    result.angle_deg = aa * 180.0 / PI;
    const double axis_angle[3] = {q[0] / target.RmsRadius(), q[1] / target.RmsRadius(), q[2] / target.RmsRadius()};
    double column[3][3];
    for (int j = 0; j < 3; j++) {
        const double e[3] = {j == 0 ? 1.0 : 0.0, j == 1 ? 1.0 : 0.0, j == 2 ? 1.0 : 0.0};
        ceres::AngleAxisRotatePoint(axis_angle, e, column[j]);
    }
    result.rotation = gemmi::Mat33(column[0][0], column[1][0], column[2][0],
                                   column[0][1], column[1][1], column[2][1],
                                   column[0][2], column[1][2], column[2][2]);
    result.shift_A = std::sqrt(q[3] * q[3] + q[4] * q[4] + q[5] * q[5]);
    result.seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();

    if (!result.zones.empty())
        logger.Info("Model validation: rigid body over {} resolution zone(s) down to {:.1f} A, "
                    "{} evaluations and {} Jacobians in {:.3f} s: rotation {:.3f} deg, translation {:.3f} A{}",
                    result.zones.size(), result.zones.back(), result.evaluations, result.jacobians,
                    result.seconds, result.angle_deg, result.shift_A, gpu != nullptr ? " (GPU)" : "");
    else
        logger.Info("Model validation: rigid body had no resolution zone with enough reflections to "
                    "run in; the model is left where it arrived");
    return result;
}