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

#include "Refinement.h"

#include <algorithm>
#include <chrono>
#include <future>
#include <map>
#include <optional>
#include <cmath>
#include <random>
#include <thread>
#include <unordered_map>

#include <ceres/gradient_problem.h>
#include <ceres/gradient_problem_solver.h>
#include <ceres/jet.h>
#include <ceres/rotation.h>

#include "gemmi/fourier.hpp"   // get_size_for_hkl, get_f_phi_on_grid
#include "gemmi/it92.hpp"
#include "gemmi/modify.hpp"    // remove_hydrogens
#include "gemmi/scaling.hpp"
#include "gemmi/stats.hpp"     // calculate_data_statistics

#include "../../common/JFJochException.h"
#include "../../common/JFJochMath.h"
#include "../../common/Logger.h"
#include "ModelFFT.h"
#include "ModelScaling.h"
#include "RefinementMath.h"
#include "RefinementMinimiser.h"
#include "RefinementXray.h"
#ifdef JFJOCH_USE_CUDA
#include "../../common/CUDAWrapper.h"
#include "RefinementGPU.h"
#endif

namespace {

constexpr double B_FLOOR = 2.0;      // A^2: no atom's B goes below it
// nor above twice the model's largest B, within these bounds (the GPU's gather is reserved for it)
constexpr double B_CEILING_MIN = 100.0, B_CEILING_MAX = 300.0;
// How far below the smallest B of the model a macro-cycle's B refinement may take an atom; the blur of the
// cycle is chosen for that.
constexpr double B_STEP_DOWN = 10.0;
constexpr double REPULSION_MARGIN = 0.5;   // A: pairs this much beyond their contact distance are listed too
// The geometry the weight is chosen within (Afonine et al. 2012 select among weights by R-free with the
// geometry held to a target in the same way).
constexpr double MAX_RMSD_BONDS = 0.02, MAX_RMSD_ANGLES = 2.0;
// R-free differences smaller than this do not decide between two weights: the tighter geometry is taken.
constexpr double R_FREE_TIE = 0.001;
// The minimiser (RefinementMinimiser): a solve stops where an iteration lowers the target by less than this,
// relative; the most a position (A) or a B (A^2) moves in the first trial of a line search.
constexpr double MINIMISER_F_TOLERANCE = 1e-7;
constexpr double MAX_STEP_XYZ = 0.3, MAX_STEP_B = 10.0;

// The figure of merit of the likelihood, <cos(dphi)>, at X = 2 Fo D Fc / var: I1(X)/I0(X) for an acentric
// reflection, tanh(X / 2) for a centric one.
// Following Murshudov, Vagin & Dodson (1997) Acta Cryst. D53, 240-255
double Fom(double x, bool centric) {
    double log_i0 = 0, ratio = 0;
    refine_math::BesselI0I1(x, log_i0, ratio);
    return centric ? std::tanh(0.5 * x) : ratio;
}

std::string Joined(const std::vector<std::string> &names) {
    std::string out;
    for (const std::string &n : names)
        out += (out.empty() ? "" : " ") + n;
    return out;
}

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

// The model as refinement takes it: the first model only, without hydrogens and atoms of unknown element,
// every B isotropic (an anisotropic one becomes its equivalent isotropic B) and inside the allowed range.
gemmi::Structure PrepareModel(const gemmi::Structure &placed, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg) {
    gemmi::Structure st = placed;
    st.models.resize(1);
    st.cell = cell;
    st.spacegroup_hm = sg.xhm();
    gemmi::remove_hydrogens(st.models[0]);
    for (gemmi::Chain &ch : st.models[0].chains)
        for (gemmi::Residue &r : ch.residues) {
            std::erase_if(r.atoms, [](const gemmi::Atom &a) { return a.element == gemmi::El::X; });
            for (gemmi::Atom &a : r.atoms) {
                if (a.aniso.nonzero()) {
                    a.b_iso = static_cast<float>(gemmi::u_to_b() * a.aniso.trace() / 3.0);
                    a.aniso = {0, 0, 0, 0, 0, 0};
                }
                a.b_iso = static_cast<float>(std::clamp<double>(a.b_iso, B_FLOOR, B_CEILING_MAX));
            }
        }
    gemmi::remove_empty_children(st);
    st.setup_cell_images();
    return st;
}

// What every stage works on: the model, the reflections, the engine and the row model of the cycle.
struct Context {
    gemmi::Structure st;
    std::vector<gemmi::Atom *> atoms;
    gemmi::UnitCell cell;
    const gemmi::SpaceGroup *sg = nullptr;
    double d_min = 0;
    RefinementReflections refl;
    std::unique_ptr<RefinementEngine> engine;
    std::unique_ptr<BulkSolventMask> mask;
#ifdef JFJOCH_USE_CUDA
    std::unique_ptr<ModelScaleGPUPool> scale_gpu;
#endif
    size_t nthreads = 1;
    // the likelihood fit's: all of the run's, also in a trial's slot - the slots refit at the same time, the GPU
    // idles meanwhile, and the bins they queue on the shared worker pool are fitted by whichever thread is free
    size_t ml_threads = 1;
    int device = -1;            // the engine's GPU
    size_t engine_bytes = 0;    // and what it holds there
    XrayRowModel rows;
    std::vector<std::complex<float>> fcalc, fmask;
    double k_sol = 0, b_sol = 0;
    // where the last Refit()'s solvent grid ended (before MoveOverallB()): the next one's fine pass is around it
    // (negative: the next one searches the whole box)
    double grid_k_sol = -1, grid_b_sol = -1;
    double k_overall = 0;
    gemmi::SMat33<double> b_star{0, 0, 0, 0, 0, 0};
    double b_floor = B_FLOOR;   // the cycle's lowest allowed B
    double b_ceiling = B_CEILING_MAX;
    // Riding hydrogens: per atom the parent it rides on (-1: none), and where it sits from the parent.
    std::vector<int> rider_parent;
    std::vector<gemmi::Position> rider_offset;
    int evaluations = 0;
    const std::atomic<bool> *cancel = nullptr;   // RefinementSettings::cancel

    bool Cancelled() const { return cancel != nullptr && cancel->load(std::memory_order_relaxed); }

    gemmi::Model &Model() { return st.models[0]; }

    // The hydrogens moved with their parents, at the offsets they were last placed at, and given their B.
    void MoveRiders() {
        for (size_t i = 0; i < rider_parent.size(); i++)
            if (rider_parent[i] >= 0) {
                atoms[i]->pos = atoms[rider_parent[i]]->pos + rider_offset[i];
                atoms[i]->b_iso = atoms[rider_parent[i]]->b_iso;
            }
    }
    // What the target says of a hydrogen is said of its parent.
    void FoldRiders(std::vector<std::array<double, 4>> &g) const {
        for (size_t i = 0; i < rider_parent.size(); i++)
            if (rider_parent[i] >= 0)
                for (int c = 0; c < 4; c++)
                    g[rider_parent[i]][c] += g[i][c];
    }

    std::vector<std::array<double, 4>> XyzB() const {
        std::vector<std::array<double, 4>> out(atoms.size());
        for (size_t i = 0; i < atoms.size(); i++)
            out[i] = {atoms[i]->pos.x, atoms[i]->pos.y, atoms[i]->pos.z, atoms[i]->b_iso};
        return out;
    }
    std::vector<gemmi::Position> Positions() const {
        std::vector<gemmi::Position> out;
        for (const gemmi::Atom *a : atoms)
            out.push_back(a->pos);
        return out;
    }
    void SetPositions(const std::vector<gemmi::Position> &p) {
        for (size_t i = 0; i < atoms.size(); i++)
            atoms[i]->pos = p[i];
    }
    std::vector<double> B() const {
        std::vector<double> out;
        for (const gemmi::Atom *a : atoms)
            out.push_back(a->b_iso);
        return out;
    }
    void SetB(const std::vector<double> &b) {
        for (size_t i = 0; i < atoms.size(); i++)
            atoms[i]->b_iso = static_cast<float>(b[i]);
    }

    // What Refit() fitted, to start another context from without fitting it again.
    struct Fit {
        double b_floor = B_FLOOR;
        double k_sol = 0, b_sol = 0, k_overall = 0;
        gemmi::SMat33<double> b_star{0, 0, 0, 0, 0, 0};
        XrayRowModel rows;
        double grid_k_sol = -1, grid_b_sol = -1;
    };
    Fit Fitted() const { return {b_floor, k_sol, b_sol, k_overall, b_star, rows, grid_k_sol, grid_b_sol}; }
    // Refit() as it was where `fit` was taken, for the atoms as they were there.
    void Restore(const Fit &fit) {
        b_floor = fit.b_floor;
        engine->SetBlur(RefinementBlur(cell, *sg, d_min, b_floor));
        k_sol = fit.k_sol;
        b_sol = fit.b_sol;
        k_overall = fit.k_overall;
        b_star = fit.b_star;
        rows = fit.rows;
        grid_k_sol = fit.grid_k_sol;
        grid_b_sol = fit.grid_b_sol;
        engine->SetRowModel(rows);
    }
    // The next Refit() from nothing an earlier one found: the whole solvent box, the likelihood from the moments.
    void ForgetStarts() {
        grid_k_sol = grid_b_sol = -1;
        rows.bin_d.clear();
    }

    // A new macro-cycle: the blur for the B the cycle may reach, the bulk-solvent mask of the model as it is,
    // the scale and the solvent fitted to the working set, and the likelihood's parameters from the free set.
    // Returns the R-factors of the model as it stands.
    RFactors Refit() {
        const RefinementStage stage("refit");
        double b_min = b_ceiling;
        for (const gemmi::Atom *a : atoms)
            b_min = std::min<double>(b_min, a->b_iso);
        b_floor = std::max(B_FLOOR, b_min - B_STEP_DOWN);
        engine->SetBlur(RefinementBlur(cell, *sg, d_min, b_floor));
        {
            const RefinementStage s1("refit.fcalc");
            engine->Fcalc(Model(), fcalc);
        }
        {
            const RefinementStage s2("refit.fmask");
            // the GPU engine masks on the same grid as BulkSolventMask (bit for bit but at points on an atom's
            // radius, ModelMaskGPU.h)
            if (engine->OnGPU())
                engine->Fmask(Model(), fmask);
            else
                mask->Fmask(Model(), refl.hkl, fmask);
        }
        std::optional<RefinementStage> s3("refit.scale");

        gemmi::AsuData<std::complex<float>> fc, fm;
        gemmi::AsuData<gemmi::ValueSigma<float>> fo;
        fc.unit_cell_ = fm.unit_cell_ = fo.unit_cell_ = cell;
        fc.spacegroup_ = fm.spacegroup_ = fo.spacegroup_ = sg;
        for (size_t m = 0; m < refl.Size(); m++) {
            fc.v.push_back({refl.hkl[m], fcalc[m]});
            fm.v.push_back({refl.hkl[m], fmask[m]});
            if (!refl.free[m])
                fo.v.push_back({refl.hkl[m], {refl.fo[m], 1.0f}});
        }
        gemmi::Scaling<float> scaling(cell, sg);
        scaling.use_solvent = true;
        scaling.prepare_points(fc, fo, &fm);
        // the scale of the score: by robust least squares (ModelScaleBox::robust)
#ifdef JFJOCH_USE_CUDA
        if (scale_gpu)
            scale_gpu->Fit(scaling, SolventBMax(Model()), nthreads, true, grid_k_sol, grid_b_sol);
        else
#endif
            FitModelScale(scaling, {.b_hi = SolventBMax(Model()), .robust = true, .k_start = grid_k_sol,
                                    .b_start = grid_b_sol}, nthreads);
        s3.reset();
        k_sol = grid_k_sol = scaling.k_sol;
        b_sol = grid_b_sol = scaling.b_sol;
        k_overall = scaling.k_overall;
        b_star = scaling.b_star;
        rows.k_total.resize(refl.Size());
        rows.fmask.resize(refl.Size());
        std::vector<std::complex<float>> fmodel(refl.Size());
        for (size_t m = 0; m < refl.Size(); m++) {
            rows.k_total[m] = static_cast<float>(scaling.get_overall_scale_factor(refl.hkl[m]));
            rows.fmask[m] = static_cast<float>(scaling.get_solvent_scale(cell.calculate_stol_sq(refl.hkl[m]))) * fmask[m];
            fmodel[m] = rows.k_total[m] * (fcalc[m] + rows.fmask[m]);
        }
        {
            const RefinementStage s4("refit.ml");
            EstimateLikelihoodParameters(refl, fmodel, rows, ml_threads);
        }
        engine->SetRowModel(rows);
        return ComputeRFactors(refl, fmodel);
    }

    // The overall isotropic B of the last Refit() moved into the atoms, and into the solvent's B so that F_model stays
    // as it is (but where an atom meets a bound of B): left in the scale, it and the atoms' B drift against each
    // other as the B factors are refined - and the B restraints, in proportion to B, depend on which holds it.
    // As phenix.refine moves the overall B into the atoms (Afonine et al. (2012) Acta Cryst. D68, 352-367)
    void MoveOverallB() {
        const gemmi::Mat33 b = cell.orth.mat.multiply(b_star.as_mat33()).multiply(cell.orth.mat.transpose());
        const double b_iso = (b[0][0] + b[1][1] + b[2][2]) / 3.0;
        double b_min = b_ceiling;
        for (size_t i = 0; i < atoms.size(); i++)
            if (rider_parent.empty() || rider_parent[i] < 0) {
                atoms[i]->b_iso = static_cast<float>(std::clamp(atoms[i]->b_iso + b_iso, B_FLOOR, b_ceiling));
                b_min = std::min<double>(b_min, atoms[i]->b_iso);
            }
        MoveRiders();
        b_floor = std::max(B_FLOOR, b_min - B_STEP_DOWN);
        engine->SetBlur(RefinementBlur(cell, *sg, d_min, b_floor));
        // b_iso I in Cartesian is b_iso F F^T in the reciprocal basis
        const gemmi::Mat33 f = cell.frac.mat.multiply(cell.frac.mat.transpose());
        b_star = {b_star.u11 - b_iso * f[0][0], b_star.u22 - b_iso * f[1][1], b_star.u33 - b_iso * f[2][2],
                  b_star.u12 - b_iso * f[0][1], b_star.u13 - b_iso * f[0][2], b_star.u23 - b_iso * f[1][2]};
        b_sol += b_iso;
        for (size_t m = 0; m < refl.Size(); m++) {
            rows.k_total[m] *= static_cast<float>(std::exp(b_iso * refl.inv_d2[m] / 4.0));
            rows.fmask[m] *= static_cast<float>(std::exp(-b_iso * refl.inv_d2[m] / 4.0));
        }
        engine->SetRowModel(rows);
    }

    // Every term of the score of the atoms as they are, at the scale, the solvent and the likelihood parameters
    // of the last Refit().
    RefinementScore Score() {
        const std::vector<std::complex<float>> fmodel = FModel();
        RefinementScore sc;
        sc.k_overall = k_overall;
        sc.k_sol = k_sol;
        sc.b_sol = b_sol;
        sc.mask_spacing = mask->Spacing();
        // b_star is in the reciprocal basis (exp(-h^T b_star h / 4)); in Cartesian B = O b_star O^T
        const gemmi::Mat33 b = cell.orth.mat.multiply(b_star.as_mat33()).multiply(cell.orth.mat.transpose());
        sc.b_iso = (b[0][0] + b[1][1] + b[2][2]) / 3.0;
        sc.b_aniso = {b[0][0] - sc.b_iso, b[1][1] - sc.b_iso, b[2][2] - sc.b_iso, b[0][1], b[0][2], b[1][2]};
        const RFactors r = ComputeRFactors(refl, fmodel);
        sc.r_work = r.r_work;
        sc.r_free = r.r_free;
        sc.n_work = r.n_work;
        sc.n_free = r.n_free;
        sc.nll_work = sc.nll_free = 0;
        double fom_sum = 0;
        int fom_n = 0;
        for (size_t m = 0; m < refl.Size(); m++) {
            double dfc = 0;
            const double fc = std::abs(fmodel[m]);
            const double nll = refine_math::AmplitudeNegLogLikelihood<double>(refl.fo[m], fc, rows.d[m], rows.var[m],
                                                                             refl.centric[m], dfc);
            (refl.free[m] ? sc.nll_free : sc.nll_work) += nll;
            // the figure of merit of the same likelihood, over the working reflections
            const double x = 2.0 * refl.fo[m] * rows.d[m] * fc / rows.var[m];
            if (!refl.free[m]) {
                fom_sum += Fom(x, refl.centric[m]);
                fom_n++;
            }
        }
        sc.fom = fom_n > 0 ? fom_sum / fom_n : NAN;
        for (int p : rider_parent)
            sc.hydrogens += p >= 0;
        constexpr int SHELLS = 20;
        double s_lo = refl.inv_d2[0], s_hi = refl.inv_d2[0];
        for (double s2 : refl.inv_d2) {
            s_lo = std::min(s_lo, s2);
            s_hi = std::max(s_hi, s2);
        }
        std::vector<std::array<double, 4>> sums(SHELLS, {0, 0, 0, 0});
        sc.shells.resize(SHELLS);
        for (int k = 0; k < SHELLS; k++) {
            sc.shells[k].d_max = 1.0 / std::sqrt(s_lo + (s_hi - s_lo) * k / SHELLS);
            sc.shells[k].d_min = 1.0 / std::sqrt(s_lo + (s_hi - s_lo) * (k + 1) / SHELLS);
        }
        for (size_t m = 0; m < refl.Size(); m++) {
            const int k = std::min(SHELLS - 1, static_cast<int>((refl.inv_d2[m] - s_lo) / (s_hi - s_lo) * SHELLS));
            const int f = refl.free[m] ? 2 : 0;
            sums[k][f] += std::fabs(refl.fo[m] - std::abs(fmodel[m]));
            sums[k][f + 1] += refl.fo[m];
            (refl.free[m] ? sc.shells[k].n_free : sc.shells[k].n_work)++;
        }
        for (int k = 0; k < SHELLS; k++) {
            if (sums[k][1] > 0)
                sc.shells[k].r_work = sums[k][0] / sums[k][1];
            if (sums[k][3] > 0)
                sc.shells[k].r_free = sums[k][2] / sums[k][3];
        }
        return sc;
    }

    // The X-ray target's curvature in each atom's B, up to one factor (CalibrateBCurvature): the likelihood's
    // curvature in |F_model| of the cycle, 2 (k_total D)^2 / var, summed in resolution shells, against the atom's own
    // part of it, (s^2 / 4)^2 f(s)^2 exp(-B s^2 / 2). A hydrogen's is its parent's.
    std::vector<double> XrayBCurvatureShape() const {
        constexpr int SHELLS = 32;
        double s2_max = 0;
        for (double s2 : refl.inv_d2)
            s2_max = std::max(s2_max, s2);
        std::vector<double> w(SHELLS, 0), s2_sum(SHELLS, 0);
        std::vector<int> n(SHELLS, 0);
        for (size_t m = 0; m < refl.Size(); m++) {
            if (refl.free[m])
                continue;
            const int k = std::min(SHELLS - 1, static_cast<int>(refl.inv_d2[m] / s2_max * SHELLS));
            const double kd = rows.k_total[m] * rows.d[m];
            w[k] += 2.0 * kd * kd / rows.var[m];
            s2_sum[k] += refl.inv_d2[m];
            n[k]++;
        }
        std::map<std::pair<int, signed char>, std::vector<double>> f2;   // per element and charge, per shell
        std::vector<double> shape(atoms.size(), 0.0);
        for (size_t i = 0; i < atoms.size(); i++) {
            const gemmi::Atom &a = *atoms[i];
            std::vector<double> &f = f2[{static_cast<int>(a.element.elem), a.charge}];
            if (f.empty()) {
                const auto &coef = gemmi::IT92<double>::get(a.element, a.charge, a.serial);
                f.resize(SHELLS, 0);
                for (int k = 0; k < SHELLS; k++)
                    if (n[k] > 0)
                        f[k] = std::pow(coef.calculate_sf(s2_sum[k] / n[k] / 4.0), 2);
            }
            for (int k = 0; k < SHELLS; k++)
                if (n[k] > 0) {
                    const double s2 = s2_sum[k] / n[k];
                    shape[i] += w[k] * f[k] * a.occ * a.occ * std::exp(-a.b_iso * s2 / 2.0) * s2 * s2 / 16.0;
                }
        }
        for (size_t i = 0; i < rider_parent.size(); i++)
            if (rider_parent[i] >= 0)
                shape[rider_parent[i]] += shape[i];
        return shape;
    }

    // The factor that turns XrayBCurvatureShape() into the X-ray target's curvature: its second derivative along a
    // random +-1 direction of the B factors of `params`, by central differences of the gradient, over what the shape
    // predicts for it.
    double CalibrateBCurvature(const std::vector<int> &params, const std::vector<double> &shape) {
        constexpr double EPS = 0.5;   // A^2
        const std::vector<double> b0 = B();
        std::mt19937 rng(7);
        std::vector<double> v(params.size());
        for (double &x : v)
            x = rng() % 2 == 0 ? 1.0 : -1.0;
        std::array<std::vector<std::array<double, 4>>, 2> g;
        for (int side = 0; side < 2; side++) {
            for (size_t k = 0; k < params.size(); k++)
                atoms[params[k]]->b_iso = static_cast<float>(b0[params[k]] + (side == 0 ? EPS : -EPS) * v[k]);
            MoveRiders();
            engine->Evaluate(Model(), &g[side]);
            FoldRiders(g[side]);
        }
        SetB(b0);
        MoveRiders();
        double num = 0, den = 0;
        for (size_t k = 0; k < params.size(); k++) {
            num += v[k] * (g[0][params[k]][3] - g[1][params[k]][3]) / (2 * EPS);
            den += shape[params[k]];
        }
        return den > 0 ? std::max(num / den, 0.0) : 0.0;
    }

    // F_model with the cycle's scale and solvent, for the atoms as they are now.
    std::vector<std::complex<float>> FModel() {
        engine->Fcalc(Model(), fcalc);
        std::vector<std::complex<float>> fmodel(refl.Size());
        for (size_t m = 0; m < refl.Size(); m++)
            fmodel[m] = rows.k_total[m] * (fcalc[m] + rows.fmask[m]);
        return fmodel;
    }
};

ceres::GradientProblemSolver::Options SolverOptions(int iterations) {
    ceres::GradientProblemSolver::Options options;
    options.line_search_direction_type = ceres::LBFGS;
    options.max_num_iterations = iterations;
    options.function_tolerance = 1e-7;
    options.gradient_tolerance = 1e-10;
    options.parameter_tolerance = 1e-9;
    options.logging_type = ceres::SILENT;
    return options;
}

// SolverOptions(iterations) that stop at the next iteration, where the minimiser stands, once ctx is cancelled.
class CancelCheck : public ceres::IterationCallback {
public:
    explicit CancelCheck(const Context &ctx) : ctx_(ctx) {}
    ceres::CallbackReturnType operator()(const ceres::IterationSummary &) override {
        return ctx_.Cancelled() ? ceres::SOLVER_TERMINATE_SUCCESSFULLY : ceres::SOLVER_CONTINUE;
    }
    ceres::GradientProblemSolver::Options Options(int iterations) const {
        ceres::GradientProblemSolver::Options options = SolverOptions(iterations);
        options.callbacks.push_back(const_cast<CancelCheck *>(this));
        return options;
    }

private:
    const Context &ctx_;
};

// The positions or the B factors of the refinable atoms, against w times the X-ray target plus a restraint set
// (minimised by RefinementMinimiser).
class AtomTarget {
public:
    enum class Mode { Xyz, B };
    AtomTarget(Context &ctx, Mode mode, const std::vector<int> &param_atoms, double weight, int restraint_set)
        : ctx_(ctx), mode_(mode), param_atoms_(param_atoms), weight_(weight), restraint_set_(restraint_set) {}

    // The parameters of the atoms as they are: x y z per atom, or B per atom.
    std::vector<double> Get() const {
        std::vector<double> p;
        for (int i : param_atoms_) {
            const gemmi::Atom &a = *ctx_.atoms[i];
            if (mode_ == Mode::Xyz) {
                p.push_back(a.pos.x);
                p.push_back(a.pos.y);
                p.push_back(a.pos.z);
            } else
                p.push_back(a.b_iso);
        }
        return p;
    }
    // The atoms set to the parameters, their riders with them.
    void Set(const double *p) const {
        for (size_t k = 0; k < param_atoms_.size(); k++) {
            gemmi::Atom &a = *ctx_.atoms[param_atoms_[k]];
            if (mode_ == Mode::Xyz)
                a.pos = gemmi::Position(p[3 * k], p[3 * k + 1], p[3 * k + 2]);
            else
                a.b_iso = static_cast<float>(std::clamp(p[k], ctx_.b_floor, ctx_.b_ceiling));
        }
        ctx_.MoveRiders();
    }
    // the most a parameter moves in the first trial of a line search
    double MaxStep() const { return mode_ == Mode::Xyz ? MAX_STEP_XYZ : MAX_STEP_B; }

    bool Evaluate(const double *p, double *cost, double *gradient) const {
        const RefinementStage stage(mode_ == Mode::Xyz ? "eval.xyz" : "eval.b");
        Set(p);
        ctx_.evaluations++;
        std::vector<std::array<double, 4>> gx, gr;
        double xray = 0, restraints = 0;
        {
            const RefinementStage s1("eval.engine");
            xray = ctx_.engine->EvaluateWithRestraints(ctx_.Model(), restraint_set_, gradient != nullptr ? &gx : nullptr,
                                                       gradient != nullptr ? &gr : nullptr, restraints);
        }
        *cost = weight_ * xray + restraints;
        if (gradient == nullptr)
            return std::isfinite(*cost);
        ctx_.FoldRiders(gx);
        for (size_t k = 0; k < param_atoms_.size(); k++) {
            const int i = param_atoms_[k];
            if (mode_ == Mode::Xyz)
                for (int c = 0; c < 3; c++)
                    gradient[3 * k + c] = weight_ * gx[i][c] + gr[i][c];
            else {
                // held at a bound, B does not move further out: by the whole target, restraints included
                const double g = weight_ * gx[i][3] + gr[i][3];
                const bool held = (p[k] <= ctx_.b_floor && g > 0) || (p[k] >= ctx_.b_ceiling && g < 0);
                gradient[k] = held ? 0.0 : g;
            }
        }
        return std::isfinite(*cost);
    }

private:
    Context &ctx_;
    Mode mode_;
    const std::vector<int> &param_atoms_;
    double weight_;
    int restraint_set_;
};

// The minimiser's diagonal for the B factors of `param_atoms` (RefinementMinimiser::SetDiagonal): weight times the
// X-ray target's curvature in B plus the B restraints', each element at least a thousandth of the median. The
// positions take none: their restraints' curvature is that of stiff bonds and angles, which move atoms together, so
// as a diagonal it slows rather than scales the minimiser.
std::vector<double> BDiagonal(const std::vector<double> &xray, const std::vector<double> &restraints,
                              const std::vector<int> &param_atoms, double weight) {
    std::vector<double> h;
    for (int i : param_atoms)
        h.push_back(weight * xray[i] + restraints[i]);
    std::vector<double> sorted = h;
    std::nth_element(sorted.begin(), sorted.begin() + sorted.size() / 2, sorted.end());
    const double floor = std::max(1e-3 * sorted[sorted.size() / 2], 1e-12);
    for (double &v : h)
        v = std::max(v, floor);
    return h;
}

// `minimiser` carries its memory from the call before; `diagonal` is the curvature it scales the steps by (none: 1).
void Minimise(Context &ctx, const AtomTarget &target, int iterations, RefinementMinimiser &minimiser,
              std::vector<double> diagonal = {}) {
    std::vector<double> p = target.Get();
    if (p.empty())
        return;
    minimiser.SetDiagonal(diagonal.empty() ? std::vector<double>(p.size(), 1.0) : std::move(diagonal));
    minimiser.Minimise([&](const std::vector<double> &x, double &f, std::vector<double> &g) {
                           return target.Evaluate(x.data(), &f, g.data());
                       },
                       p, iterations, MINIMISER_F_TOLERANCE, target.MaxStep(), [&]() { return ctx.Cancelled(); });
    target.Set(p.data());   // where the minimiser ended, not where its last line search probed
}

// Groups of atoms that move as rigid bodies (one per chain) against the X-ray target: per body an angle-axis
// rotation about its centroid, scaled by its rms radius so that a unit moves a typical atom by an angstrom as a
// unit of translation does, then a translation.
class RigidBodiesTarget : public ceres::FirstOrderFunction {
public:
    RigidBodiesTarget(Context &ctx, const std::vector<std::vector<int>> &bodies)
        : ctx_(ctx), bodies_(bodies), base_(ctx.Positions()) {
        for (const std::vector<int> &b : bodies) {
            gemmi::Position c;
            for (int i : b)
                c += base_[i];
            c *= 1.0 / static_cast<double>(b.size());
            double r2 = 0;
            for (int i : b)
                r2 += c.dist_sq(base_[i]);
            centre_.push_back(c);
            radius_.push_back(std::max(1.0, std::sqrt(r2 / static_cast<double>(b.size()))));
        }
    }
    int NumParameters() const override { return static_cast<int>(6 * bodies_.size()); }

    bool Evaluate(const double *q, double *cost, double *gradient) const override {
        const RefinementStage stage("eval.rigid");
        Place(q);
        ctx_.evaluations++;
        std::vector<std::array<double, 4>> g;
        *cost = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &g : nullptr);
        if (gradient == nullptr)
            return std::isfinite(*cost);
        using Jet = ceres::Jet<double, 3>;
        for (size_t b = 0; b < bodies_.size(); b++) {
            const Jet aa[3] = {Jet(q[6 * b] / radius_[b], 0), Jet(q[6 * b + 1] / radius_[b], 1),
                               Jet(q[6 * b + 2] / radius_[b], 2)};
            double dq[6] = {0, 0, 0, 0, 0, 0};
            for (int i : bodies_[b]) {
                const gemmi::Position d = base_[i] - centre_[b];
                const Jet p[3] = {Jet(d.x), Jet(d.y), Jet(d.z)};
                Jet rp[3];
                ceres::AngleAxisRotatePoint(aa, p, rp);
                for (int c = 0; c < 3; c++) {
                    for (int k = 0; k < 3; k++)
                        dq[k] += g[i][c] * rp[c].v[k] / radius_[b];
                    dq[3 + c] += g[i][c];
                }
            }
            std::copy(dq, dq + 6, gradient + 6 * b);
        }
        return std::isfinite(*cost);
    }

    void Place(const double *q) const {
        for (size_t b = 0; b < bodies_.size(); b++) {
            const double aa[3] = {q[6 * b] / radius_[b], q[6 * b + 1] / radius_[b], q[6 * b + 2] / radius_[b]};
            for (int i : bodies_[b]) {
                const gemmi::Position d = base_[i] - centre_[b];
                const double p[3] = {d.x, d.y, d.z};
                double rp[3];
                ceres::AngleAxisRotatePoint(aa, p, rp);
                ctx_.atoms[i]->pos = gemmi::Position(rp[0] + centre_[b].x + q[6 * b + 3], rp[1] + centre_[b].y + q[6 * b + 4],
                                                     rp[2] + centre_[b].z + q[6 * b + 5]);
            }
        }
    }

private:
    Context &ctx_;
    const std::vector<std::vector<int>> &bodies_;
    std::vector<gemmi::Position> base_;
    std::vector<gemmi::Position> centre_;
    std::vector<double> radius_;
};

// One B shift per group of atoms (a residue), against the X-ray target.
class GroupBTarget : public ceres::FirstOrderFunction {
public:
    GroupBTarget(Context &ctx, const std::vector<std::vector<int>> &groups)
        : ctx_(ctx), groups_(groups), base_(ctx.B()) {}
    int NumParameters() const override { return static_cast<int>(groups_.size()); }

    bool Evaluate(const double *q, double *cost, double *gradient) const override {
        const RefinementStage stage("eval.group_b");
        Place(q);
        ctx_.evaluations++;
        std::vector<std::array<double, 4>> g;
        *cost = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &g : nullptr);
        if (gradient != nullptr)
            for (size_t k = 0; k < groups_.size(); k++) {
                gradient[k] = 0;
                for (int i : groups_[k]) {
                    const double b = base_[i] + q[k];
                    if (b > ctx_.b_floor && b < ctx_.b_ceiling)
                        gradient[k] += g[i][3];
                }
            }
        return std::isfinite(*cost);
    }

    void Place(const double *q) const {
        for (size_t k = 0; k < groups_.size(); k++)
            for (int i : groups_[k])
                ctx_.atoms[i]->b_iso = static_cast<float>(std::clamp(base_[i] + q[k], ctx_.b_floor, ctx_.b_ceiling));
    }

private:
    Context &ctx_;
    const std::vector<std::vector<int>> &groups_;
    std::vector<double> base_;
};

// The placement stage: rigid bodies, then group B.
void RefinePlacement(Context &ctx, bool group_b, RefinementResult &result) {
    const RefinementStage stage("placement");
    std::vector<std::vector<int>> bodies, groups;
    size_t i = 0;
    for (const gemmi::Chain &ch : ctx.Model().chains) {
        std::vector<int> body;
        for (const gemmi::Residue &r : ch.residues) {
            std::vector<int> group;
            for (size_t k = 0; k < r.atoms.size(); k++, i++) {
                group.push_back(static_cast<int>(i));
                if (!r.is_water())
                    body.push_back(static_cast<int>(i));
            }
            groups.push_back(group);
        }
        // A chain of a few atoms (an ion, a ligand of its own chain) stays where it is.
        if (body.size() >= 50)
            bodies.push_back(body);
    }
    // Each step is kept only where it lowers R-free beyond the tie: a deposited model was refined atom by atom
    // already, and moving its chains or shifting its B factors residue by residue then only makes it worse.
    RFactors before = ctx.Refit();
    if (!bodies.empty()) {
        const std::vector<gemmi::Position> pos0 = ctx.Positions();
        std::vector<double> q(6 * bodies.size(), 0.0);
        auto target = std::make_unique<RigidBodiesTarget>(ctx, bodies);
        RigidBodiesTarget *rb = target.get();
        ceres::GradientProblem problem(std::move(target));
        ceres::GradientProblemSolver::Summary summary;
        const CancelCheck cancel_check(ctx);
        ceres::Solve(cancel_check.Options(30), problem, q.data(), &summary);
        rb->Place(q.data());
        const RFactors after = ctx.Refit();
        if (after.r_free < before.r_free - R_FREE_TIE) {
            result.rigid_bodies = static_cast<int>(bodies.size());
            before = after;
        } else {
            ctx.SetPositions(pos0);
            before = ctx.Refit();
        }
    }
    if (group_b) {
        const std::vector<double> b0 = ctx.B();
        std::vector<double> q(groups.size(), 0.0);
        auto target = std::make_unique<GroupBTarget>(ctx, groups);
        GroupBTarget *gb = target.get();
        ceres::GradientProblem problem(std::move(target));
        ceres::GradientProblemSolver::Summary summary;
        const CancelCheck cancel_check(ctx);
        ceres::Solve(cancel_check.Options(30), problem, q.data(), &summary);
        gb->Place(q.data());
        result.group_b = ctx.Refit().r_free < before.r_free - R_FREE_TIE;
        if (!result.group_b)
            ctx.SetB(b0);
    }
}

// The X-ray weight that balances the two gradients: the norm of the geometry's gradient (taken on the model
// shaken by 0.1 A, where a refined model's geometry would have none) over the norm of the X-ray target's.
// Following Brunger, Karplus & Petsko (1989) Acta Cryst. A45, 50-61, as phenix.refine starts its weight search
// (Afonine et al. (2012) Acta Cryst. D68, 352-367)
double BalancedWeight(Context &ctx, const std::vector<int> &refinable) {
    std::vector<std::array<double, 4>> gx, gr;
    ctx.engine->Evaluate(ctx.Model(), &gx);
    std::vector<std::array<double, 4>> xyzb = ctx.XyzB();
    std::mt19937 rng(1);
    std::normal_distribution<double> shake(0.0, 0.1);
    for (std::array<double, 4> &a : xyzb)
        for (int c = 0; c < 3; c++)
            a[c] += shake(rng);
    ctx.engine->EvaluateRestraints(0, xyzb, &gr);
    double nx = 0, ng = 0;
    for (int i : refinable)
        for (int c = 0; c < 3; c++) {
            nx += gx[i][c] * gx[i][c];
            ng += gr[i][c] * gr[i][c];
        }
    return nx > 0 ? std::sqrt(ng / nx) : 1.0;
}

// The B factors' X-ray weight that balances the two gradients, as BalancedWeight() does for the positions: the B
// similarity's gradient on the model with its B shaken by 20% (a flat B, as a model may come, gives none) over the
// X-ray target's.
double BalancedADPWeight(Context &ctx, const RestraintSet &b_restraints, const std::vector<int> &refinable) {
    std::vector<std::array<double, 4>> gx;
    ctx.engine->Evaluate(ctx.Model(), &gx);
    ctx.FoldRiders(gx);
    std::vector<std::array<double, 4>> xyzb = ctx.XyzB();
    std::mt19937 rng(2);
    std::normal_distribution<double> shake(0.0, 0.2);
    for (std::array<double, 4> &a : xyzb)
        a[3] *= std::exp(shake(rng));
    std::vector<std::array<double, 4>> gr(xyzb.size(), {0, 0, 0, 0});
    EvaluateRestraints(b_restraints, xyzb, &gr);
    double nx = 0, ng = 0;
    for (int i : refinable) {
        nx += gx[i][3] * gx[i][3];
        ng += gr[i][3] * gr[i][3];
    }
    return nx > 0 ? std::sqrt(ng / nx) : 1.0;
}

// The model's map coefficients and its difference map, weighted by the refinement's own likelihood: m its figure
// of merit and D its D (2mFo - DFc, mFo - DFc).
// Following Read (1986) Acta Cryst. A42, 140-149
void MakeMaps(Context &ctx, const std::vector<MergedReflection> &merged, double wavelength_A, RefinementResult &result) {
    const RefinementStage stage("maps");
    ctx.Refit();
    const std::vector<std::complex<float>> fmodel = ctx.FModel();
    const gemmi::GroupOps gops = ctx.sg->operations();
    std::vector<double> fom(ctx.refl.Size());
    double fom_sum = 0;
    int fom_n = 0;
    for (size_t m = 0; m < ctx.refl.Size(); m++) {
        const double x = 2.0 * ctx.refl.fo[m] * ctx.rows.d[m] * std::abs(fmodel[m]) / ctx.rows.var[m];
        fom[m] = Fom(x, ctx.refl.centric[m]);
        if (!ctx.refl.free[m]) {
            fom_sum += fom[m];
            fom_n++;
        }
    }
    result.mean_fom = fom_n > 0 ? fom_sum / fom_n : 0;

    auto maps = std::make_shared<ModelMaps>();
    MtzTable &mtz = maps->coefficients;
    mtz.space_group_name = ctx.sg->xhm();
    mtz.unit_cell = UnitCell{static_cast<float>(ctx.cell.a), static_cast<float>(ctx.cell.b), static_cast<float>(ctx.cell.c),
                             static_cast<float>(ctx.cell.alpha), static_cast<float>(ctx.cell.beta),
                             static_cast<float>(ctx.cell.gamma)};
    mtz.datasets = {{"HKL_base", "HKL_base", "HKL_base", 0.0f},
                    {"refinement", "refinement", "refinement", static_cast<float>(wavelength_A)}};
    const std::pair<const char *, char> columns[] = {{"H", 'H'}, {"K", 'H'}, {"L", 'H'}, {"FP", 'F'}, {"FC", 'F'},
                                                     {"PHIC", 'P'}, {"FWT", 'F'}, {"PHWT", 'P'}, {"DELFWT", 'F'},
                                                     {"PHDELWT", 'P'}, {"FOM", 'W'}, {"FREE", 'I'}};
    for (const auto &[label, type] : columns)
        mtz.columns.push_back({label, type, type == 'H' ? 0 : 1, {}});
    gemmi::AsuData<std::complex<float>> diff;
    diff.unit_cell_ = ctx.cell;
    diff.spacegroup_ = ctx.sg;
    for (size_t m = 0; m < ctx.refl.Size(); m++) {
        const double fo = ctx.refl.fo[m], fc = std::abs(fmodel[m]);
        const double mw = fom[m], d = ctx.rows.d[m];
        const double fwt = ctx.refl.centric[m] ? mw * fo : 2 * mw * fo - d * fc;
        const double delfwt = mw * fo - d * fc;
        const double phi = std::arg(fmodel[m]);
        double phi_deg = phi * 180.0 / PI;
        if (phi_deg < 0)
            phi_deg += 360.0;
        const float row[] = {static_cast<float>(ctx.refl.hkl[m][0]), static_cast<float>(ctx.refl.hkl[m][1]),
                             static_cast<float>(ctx.refl.hkl[m][2]), static_cast<float>(fo), static_cast<float>(fc),
                             static_cast<float>(phi_deg), static_cast<float>(fwt), static_cast<float>(phi_deg),
                             static_cast<float>(delfwt), static_cast<float>(phi_deg), static_cast<float>(mw),
                             ctx.refl.free[m] ? 0.0f : 1.0f};
        for (size_t c = 0; c < mtz.columns.size(); c++)
            mtz.columns[c].values.push_back(row[c]);
        diff.v.push_back({ctx.refl.hkl[m], std::polar(static_cast<float>(delfwt), static_cast<float>(phi))});
    }
    result.maps = maps;
    // The difference map beside the anomalous one below: neither reads the other.
    std::future<void> difference_map = std::async(std::launch::async, [&diff, &result] {
        const std::array<int, 3> size = gemmi::get_size_for_hkl(diff, {{0, 0, 0}}, 3.0);
        result.difference_map = std::make_shared<gemmi::Grid<float>>(MapFromFPhi(gemmi::get_f_phi_on_grid<float>(diff, size, true)));
        result.difference_map_rms = gemmi::calculate_data_statistics(result.difference_map->data).rms;
    });

    // The anomalous difference map: F(+) - F(-) on the model's phase turned back by 90 degrees, read on the +
    // rows of an anomalous merge only (as model validation reads it).
    // Following ANODE, Thorn & Sheldrick (2011) J. Appl. Cryst. 44, 1285-1287
    const gemmi::ReciprocalAsu asu(ctx.sg);
    std::map<gemmi::Miller, float> danom;
    for (const MergedReflection &r : merged) {
        if (!std::isfinite(r.F_plus) || !std::isfinite(r.F_minus))
            continue;
        const gemmi::Miller h{{r.h, r.k, r.l}};
        if (gops.is_reflection_centric(h))
            continue;
        const auto [hasu, plus] = asu.to_asu_sign(h, gops);
        if (plus)
            danom[hasu] = r.F_plus - r.F_minus;
    }
    if (!danom.empty()) {
        gemmi::AsuData<std::complex<float>> anom;
        anom.unit_cell_ = ctx.cell;
        anom.spacegroup_ = ctx.sg;
        for (size_t m = 0; m < ctx.refl.Size(); m++) {
            const auto it = danom.find(ctx.refl.hkl[m]);
            if (it != danom.end())
                anom.v.push_back({ctx.refl.hkl[m], it->second * std::polar(1.0f, static_cast<float>(std::arg(fmodel[m]) - PI / 2))});
        }
        if (!anom.v.empty()) {
            result.anomalous_map = std::make_shared<gemmi::Grid<float>>(
                MapFromFPhi(gemmi::get_f_phi_on_grid<float>(anom, gemmi::get_size_for_hkl(anom, {{0, 0, 0}}, 3.0), true)));
            result.anomalous_map_rms = gemmi::calculate_data_statistics(result.anomalous_map->data).rms;
        }
    }
    difference_map.get();
}

// The engine for the model as it is now: on the GPU where it is wanted and takes the case within its budget,
// on the CPU otherwise.
void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logger &logger) {
    const RefinementStage stage("engine");
    ctx.engine.reset();
#ifdef JFJOCH_USE_CUDA
    ctx.scale_gpu.reset();
    if (gpu && get_gpu_count() > 0) {
        try {
            const int device = get_gpu();
            const size_t card = get_gpu_total_memory(device);
            const size_t usable = card / 5 * 4;
            const size_t budget = usable > settings.gpu_resident_bytes ? usable - settings.gpu_resident_bytes : 0;
            size_t bytes = 0;
            ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(),
                                                 ctx.b_ceiling, budget, &bytes, settings.nthreads);
            ctx.device = device;
            ctx.engine_bytes = bytes;
            if (ctx.engine)
                ctx.scale_gpu = std::make_unique<ModelScaleGPUPool>(device, ctx.refl.Size(), 1);
            else
                logger.Info("Refinement: on the CPU - the GPU cannot take this case or {:.1f} GB is over its budget",
                            bytes / 1e9);
        } catch (const JFJochException &e) {
            logger.Warning("Refinement: the GPU failed ({}), on the CPU", e.what());
            ctx.engine.reset();
            ctx.scale_gpu.reset();
        }
    }
#else
    (void) gpu;
    (void) logger;
#endif
    if (!ctx.engine)
        ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, settings.nthreads);
}

// The riding hydrogens placed on their parents as they are, and the offsets and B they ride with.
void PlaceRiders(Context &ctx, RefinementTopology &topo) {
    const RefinementStage stage("riders");
    PlaceRidingHydrogens(topo);
    ctx.rider_offset.assign(ctx.atoms.size(), gemmi::Position());
    for (size_t i = 0; i < ctx.atoms.size(); i++)
        if (ctx.rider_parent[i] >= 0) {
            ctx.rider_offset[i] = ctx.atoms[i]->pos - ctx.atoms[ctx.rider_parent[i]]->pos;
            ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso;
        }
}

// The geometry and the repulsions of the model as it is.
RestraintSet RestraintsNow(Context &ctx, const RefinementTopology &topo) {
    const RefinementStage stage("repulsions");
    RestraintSet set = topo.geometry;
    set.Append(topo.metal_bonds);
    set.Append(BuildRepulsions(topo, ctx.Model(), ctx.cell, *ctx.sg, REPULSION_MARGIN, ctx.nthreads));
    set.Finalize(ctx.atoms.size());
    return set;
}

// The B similarity of the model as it is: bonded and angle-related pairs, and atoms in contact.
RestraintSet BRestraintsNow(Context &ctx, const RefinementTopology &topo) {
    const RefinementStage stage("b_restraints");
    RestraintSet set = topo.adp;
    set.Append(BuildBContacts(topo, ctx.Model(), ctx.cell, *ctx.sg, ctx.nthreads));
    set.Finalize(ctx.atoms.size());
    return set;
}

// Where a weight trial runs: a model with its topology and its engine.
struct TrialSlot {
    Context *ctx;
    RefinementTopology *topo;
};

// A slot of its own beside the main one.
struct OwnSlot {
    Context ctx;
    RefinementTopology topo;
};

// A copy of the model as it was before the topology was made (the same topology comes out of it, hydrogens
// included), with an engine of the main one's kind - a GPU's on `device`. Null where the card cannot make one.
std::unique_ptr<OwnSlot> MakeOwnSlot(const gemmi::Structure &before_topology, const Context &main, int device,
                                     size_t nthreads, const RefinementSettings &settings) {
    auto slot = std::make_unique<OwnSlot>();
    Context &ctx = slot->ctx;
    ctx.st = before_topology;
    ctx.cell = main.cell;
    ctx.sg = main.sg;
    ctx.d_min = main.d_min;
    ctx.refl = main.refl;
    ctx.nthreads = nthreads;
    ctx.ml_threads = settings.nthreads;
    ctx.cancel = main.cancel;
    ctx.b_ceiling = main.b_ceiling;
    ctx.mask = std::make_unique<BulkSolventMask>(ctx.cell, *ctx.sg, ctx.d_min, nthreads);
    slot->topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
    ctx.atoms = AtomList(ctx.Model());
    ctx.rider_parent = slot->topo.rider_parent;
#ifdef JFJOCH_USE_CUDA
    if (main.engine->OnGPU()) {
        try {
            ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(),
                                                 ctx.b_ceiling, SIZE_MAX, nullptr, nthreads);
            if (!ctx.engine)
                return nullptr;
            ctx.scale_gpu = std::make_unique<ModelScaleGPUPool>(device, ctx.refl.Size(), 1);
        } catch (const std::exception &) {
            return nullptr;
        }
        ctx.device = device;
    }
#else
    (void) device;
#endif
    if (!ctx.engine)
        ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, nthreads);
    ctx.engine->SetRestraints(1, slot->topo.adp);
    return slot;
}

// The devices of the slots beside the main one, `wanted` of them at most: the GPUs in turn from the main
// engine's, each taking engines while they fit in its budget. -1 for each one on the CPU, where the main one and
// those beside it are no more than the threads.
std::vector<int> OwnSlotDevices(const Context &main, int wanted, const RefinementSettings &settings) {
#ifdef JFJOCH_USE_CUDA
    if (main.engine->OnGPU()) {
        const int n = get_gpu_count();
        // a slot's engine and its scale pool, the robust fit's buffers included
        const size_t slot_bytes = main.engine_bytes + ModelScaleGPUPool::DeviceBytes(main.refl.Size(), 1, true);
        std::vector<size_t> used(n, 0);
        used[main.device] = slot_bytes;
        std::vector<int> devices;
        for (int k = 1; k <= wanted; k++) {
            const int d = (main.device + k) % n;
            const size_t usable = get_gpu_total_memory(d) / 5 * 4;
            const size_t budget = usable > settings.gpu_resident_bytes ? usable - settings.gpu_resident_bytes : 0;
            if (used[d] + slot_bytes <= budget) {
                used[d] += slot_bytes;
                devices.push_back(d);
            }
        }
        return devices;
    }
#endif
    // on the CPU, no more trials at once than threads: each slot is a model, a topology and an engine of its own
    const int beside = std::min(wanted, static_cast<int>(settings.nthreads) - 1);
    return std::vector<int>(std::max(beside, 0), -1);
}

// One weight trial: `cycles` macro-cycles from pos0 and b0, of the positions at X-ray weight w and, where adp_w is
// not 0, the B factors at X-ray weight adp_w. `fit`: what was fitted at pos0 and b0 already, null to fit it here.
struct TrialOutcome {
    RefinementWeightTrial trial;
    std::vector<gemmi::Position> pos;
    std::vector<double> b;
    Context::Fit fit;   // the scale, the solvent and the likelihood's parameters fitted where it ended
    std::string line;   // what the log says of it
    bool ran = false;   // false: left out, the refinement cancelled before it
};

TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const std::vector<gemmi::Position> &pos0,
                      const std::vector<double> &b0, const Context::Fit *fit, double w, double adp_w, int cycles,
                      const RefinementSettings &settings) {
    Context &ctx = *slot.ctx;
    RefinementTopology &topo = *slot.topo;
    const RefinementStage stage(adp_w > 0 ? "rung.b" : "rung.xyz");
    ctx.SetPositions(pos0);
    ctx.SetB(b0);
    // the positions and the B factors in turn, each with a memory of its own kept from macro-cycle to macro-cycle
    const AtomTarget xyz(ctx, AtomTarget::Mode::Xyz, refinable, w, 0), b(ctx, AtomTarget::Mode::B, refinable, adp_w, 1);
    RefinementMinimiser xyz_minimiser, b_minimiser;
    double xray_scale = 0;
    for (int cycle = 0; cycle < cycles && !ctx.Cancelled(); cycle++) {
        PlaceRiders(ctx, topo);
        if (cycle == 0 && fit != nullptr)
            ctx.Restore(*fit);
        else {
            // a trial starts from nothing the slot's previous trial left
            if (cycle == 0)
                ctx.ForgetStarts();
            ctx.Refit();
        }
        if (adp_w > 0)
            ctx.MoveOverallB();
        {
            const RestraintSet set = RestraintsNow(ctx, topo);
            const RefinementStage s1("set_restraints");
            ctx.engine->SetRestraints(0, set);
        }
        {
            const RefinementStage s2("minimise.xyz");
            Minimise(ctx, xyz, settings.xyz_iterations, xyz_minimiser);
        }
        if (adp_w > 0) {
            const RestraintSet b_set = BRestraintsNow(ctx, topo);
            ctx.engine->SetRestraints(1, b_set);
            std::vector<double> xray = ctx.XrayBCurvatureShape();
            if (cycle == 0)
                xray_scale = ctx.CalibrateBCurvature(refinable, xray);
            for (double &x : xray)
                x *= xray_scale;
            const RefinementStage s3("minimise.b");
            Minimise(ctx, b, settings.b_iterations, b_minimiser,
                     BDiagonal(xray, BRestraintCurvature(b_set, ctx.XyzB()), refinable, adp_w));
        }
    }
    PlaceRiders(ctx, topo);
    TrialOutcome out;
    RefinementWeightTrial &t = out.trial;
    t.weight = w;
    t.adp_weight = adp_w;
    const RFactors rf = ctx.Refit();
    t.r_work = rf.r_work;
    t.r_free = rf.r_free;
    {
        const RefinementStage s1("measure");
        t.geometry = MeasureGeometry(topo.geometry,
                                     BuildRepulsions(topo, ctx.Model(), ctx.cell, *ctx.sg, 0.0, ctx.nthreads), ctx.XyzB(),
                                     &topo.metal_bonds);
    }
    t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES;
    out.line = fmt::format("X-ray weight {:.3g}, B {}: R-work {:.4f}, R-free {:.4f}, rmsd bonds {:.4f} A, "
                           "angles {:.2f} deg{}{}",
                           w, adp_w > 0 ? fmt::format("refined at weight {:.3g}", adp_w) : std::string("held"),
                           rf.r_work, rf.r_free, t.geometry.rmsd_bonds, t.geometry.rmsd_angles,
                           t.geometry.metal_bonds > 0 ? fmt::format(", metal bonds {:.3f} A", t.geometry.rmsd_metal)
                                                      : std::string(),
                           t.within_bounds ? "" : " (geometry out of bounds)");
    out.pos = ctx.Positions();
    out.b = ctx.B();
    out.fit = ctx.Fitted();
    out.ran = true;
    return out;
}

// The trials (w, adp_w) side by side: trial i in slot i % slots, each slot's trials in turn on a thread of its
// own (on its GPU), returned in the order given. Every trial starts from pos0 and b0 and takes nothing from the
// one before it in its slot, so the outcome is the same for any number of slots.
std::vector<TrialOutcome> RunTrials(const std::vector<TrialSlot> &trial_slots,
                                    const std::vector<std::pair<double, double>> &trials,
                                    const std::vector<int> &refinable, const std::vector<gemmi::Position> &pos0,
                                    const std::vector<double> &b0, const Context::Fit *fit, int cycles,
                                    const RefinementSettings &settings) {
    std::vector<TrialOutcome> out(trials.size());
    std::vector<std::exception_ptr> errors(trial_slots.size());
    // the threads shared among the slots that run (the outcome does not depend on how many each has)
    const size_t running = std::min(trial_slots.size(), trials.size());
    for (size_t s = 0; s < running; s++)
        trial_slots[s].ctx->nthreads = std::max<size_t>(1, settings.nthreads / running);
    auto run_slot = [&](size_t s) {
        try {
#ifdef JFJOCH_USE_CUDA
            const ScopedGpuPin pin(trial_slots[s].ctx->device >= 0 ? trial_slots[s].ctx->device : get_gpu());
#endif
            for (size_t i = s; i < trials.size(); i += trial_slots.size())
                if (!trial_slots[s].ctx->Cancelled())
                    out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, fit, trials[i].first, trials[i].second,
                                      cycles, settings);
        } catch (...) {
            errors[s] = std::current_exception();
        }
    };
    std::vector<std::thread> threads;
    for (size_t s = 1; s < std::min(trial_slots.size(), trials.size()); s++)
        threads.emplace_back(run_slot, s);
    run_slot(0);
    for (std::thread &t : threads)
        t.join();
    for (const std::exception_ptr &e : errors)
        if (e)
            std::rethrow_exception(e);
    return out;
}

}  // namespace

RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<MergedReflection> &merged,
                             const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double wavelength_A,
                             Logger &logger, const RefinementSettings &settings) {
    const auto start = std::chrono::steady_clock::now();
    RefinementResult result;
    Context ctx;
    ctx.st = PrepareModel(placed, cell, sg);
    ctx.cell = cell;
    ctx.sg = &sg;
    ctx.nthreads = settings.nthreads;
    ctx.ml_threads = settings.nthreads;
    ctx.cancel = settings.cancel;
    ctx.atoms = AtomList(ctx.Model());
    result.atoms = static_cast<int>(ctx.atoms.size());
    double b_max = 0;
    for (const gemmi::Atom *a : ctx.atoms)
        b_max = std::max<double>(b_max, a->b_iso);
    ctx.b_ceiling = std::clamp(2 * b_max, B_CEILING_MIN, B_CEILING_MAX);
    if (ctx.atoms.empty()) {
        result.failure_reason = "the model has no atoms";
        return result;
    }
    double max_inv_d2 = 0;
    for (const MergedReflection &r : merged)
        if (std::isfinite(r.F) && r.d > 0)
            max_inv_d2 = std::max(max_inv_d2, 1.0 / (double(r.d) * r.d));
    if (max_inv_d2 <= 0) {
        result.failure_reason = "no amplitudes";
        return result;
    }
    // just past the finest reflection, so that prepare_asu_data()'s strict limit keeps it
    ctx.d_min = 0.999 / std::sqrt(max_inv_d2);
    result.d_min = 1.0 / std::sqrt(max_inv_d2);
    ctx.refl = PrepareRefinementReflections(merged, cell, sg, ctx.d_min);
    if (ctx.refl.Size() < 100) {
        result.failure_reason = "too few reflections";
        return result;
    }

    ctx.mask = std::make_unique<BulkSolventMask>(cell, sg, ctx.d_min, settings.nthreads);

    // The dictionary's view of the model, and its hydrogens riding on their parents from the start, so that the
    // model is scored as the restrained stage refines it.
    const bool restrained = settings.restrained || settings.score_only;
    // The weight trials' slots start from a copy of the model as it is before the topology is made.
    const gemmi::Structure before_topology = ctx.st;
    RefinementTopology topo;
    if (restrained) {
        std::optional<RefinementStage> topo_stage("topology");
        topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
        topo_stage.reset();
        ctx.atoms = AtomList(ctx.Model());
        ctx.rider_parent = topo.rider_parent;
        PlaceRiders(ctx, topo);
    }

    MakeEngine(ctx, settings.gpu, settings, logger);
    result.on_gpu = ctx.engine->OnGPU();

    RFactors r = ctx.Refit();
    result.n_work = r.n_work;
    result.n_free = r.n_free;
    result.r_work_start = r.r_work;
    result.r_free_start = r.r_free;
    result.score_start = ctx.Score();
    logger.Info("Refinement: {} atoms ({} riding hydrogens), {} reflections to {:.2f} A ({} free), on the {}; as given "
                "R-work {:.4f}, R-free {:.4f}; k_sol {:.3f}, b_sol {:.1f}, overall B {:.1f}, mask grid {:.2f} A",
                ctx.atoms.size(), result.score_start.hydrogens, ctx.refl.Size(), result.d_min, r.n_free,
                result.on_gpu ? "GPU" : "CPU", r.r_work, r.r_free, ctx.k_sol, ctx.b_sol, result.score_start.b_iso,
                result.score_start.mask_spacing);

    // The maps, the blobs and the model of where the refinement ended.
    auto finish = [&]() {
        MakeMaps(ctx, merged, wavelength_A, result);
        {
            const RefinementStage stage("blobs");
            result.blobs = FindDensityBlobs(*result.difference_map, result.difference_map_rms, ctx.st, result.d_min,
                                            result.anomalous_map.get(), result.anomalous_map_rms);
        }
        result.k_sol = ctx.k_sol;
        result.b_sol = ctx.b_sol;
        result.model = std::make_shared<gemmi::Structure>(ctx.st);
        result.ok = true;
        result.seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
    };
    if (settings.score_only) {
        result.score_start.hkl = ctx.refl.hkl;
        result.score_start.fcalc = ctx.fcalc;
        result.score_start.fmask = ctx.rows.fmask;
        result.score_start.k_total = ctx.rows.k_total;
        result.score_start.d = ctx.rows.d;
        result.score_start.var = ctx.rows.var;
        result.r_work_placement = result.r_work_start;
        result.r_free_placement = result.r_free_start;
        result.restrained_residues = topo.restrained_residues;
        result.fixed_residues = topo.fixed_residues;
        result.fixed_names = topo.fixed_names;
        result.geometry_start = MeasureGeometry(topo.geometry, BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0),
                                                ctx.XyzB(), &topo.metal_bonds);
        finish();
        const GeometryStatistics &g = result.geometry_start;
        logger.Info("Refinement: the model scored only; geometry rmsd bonds {:.4f} A (rmsZ {:.2f}, {} bonds), metal "
                    "bonds {:.4f} A (rmsZ {:.2f}, {}; {} links made), both {:.4f} A; angles {:.2f} deg (rmsZ {:.2f})",
                    g.rmsd_bonds, g.rmsz_bonds, g.bonds, g.rmsd_metal, g.rmsz_metal, g.metal_bonds, topo.metal_links,
                    g.RmsdAllBonds(), g.rmsd_angles, g.rmsz_angles);
        return result;
    }

    // The trials run side by side, as many at a time as the longest set of them has and the cards take engines:
    // on the main model and engine, and on copies with engines of their own; each with a share of the threads. The
    // copies are made while the placement runs: they start from the model as it was before the topology was made
    // and take nothing from the main context that the placement changes.
    int wanted = static_cast<int>(std::max({settings.weight_ladder.size() * settings.adp_ladder.size(),
                                            settings.final_ladder.size(), size_t{3}}));
    if (settings.max_trials_at_once > 0)
        wanted = std::min(wanted, static_cast<int>(settings.max_trials_at_once));
    const std::vector<int> devices = settings.restrained ? OwnSlotDevices(ctx, wanted - 1, settings) : std::vector<int>();
    const size_t slot_threads = std::max<size_t>(1, settings.nthreads / (devices.size() + 1));
    std::vector<std::unique_ptr<OwnSlot>> own(devices.size());
    std::vector<std::thread> slot_builders;
    // joined before the slots are taken, or here where the placement throws
    struct JoinAll {
        std::vector<std::thread> &threads;
        ~JoinAll() {
            for (std::thread &t : threads)
                if (t.joinable())
                    t.join();
        }
    } join_slot_builders{slot_builders};
    for (size_t k = 0; k < devices.size(); k++)
        slot_builders.emplace_back([&, k]() {
#ifdef JFJOCH_USE_CUDA
            const ScopedGpuPin pin(devices[k] >= 0 ? devices[k] : get_gpu());
#endif
            // a slot that cannot be made is left out, as where the card cannot take its engine
            try {
                own[k] = MakeOwnSlot(before_topology, ctx, devices[k], slot_threads, settings);
            } catch (const std::exception &e) {
                logger.Warning("Refinement: a trial slot could not be made ({}), the trials share fewer", e.what());
            }
        });

    RefinePlacement(ctx, settings.group_b, result);
    if (restrained)
        PlaceRiders(ctx, topo);
    r = ctx.Refit();
    result.r_work_placement = r.r_work;
    result.r_free_placement = r.r_free;
    logger.Info("Refinement, placement ({} rigid bodies kept, B per residue {}): R-work {:.4f}, R-free {:.4f}",
                result.rigid_bodies, result.group_b ? "kept" : "not kept", r.r_work, r.r_free);

    if (settings.restrained) {
        ctx.Refit();
        result.restrained_residues = topo.restrained_residues;
        result.fixed_residues = topo.fixed_residues;
        result.fixed_names = topo.fixed_names;
        std::vector<int> refinable;
        for (size_t i = 0; i < ctx.atoms.size(); i++)
            if (topo.refinable[i])
                refinable.push_back(static_cast<int>(i));
        result.refined_atoms = static_cast<int>(refinable.size());
        logger.Info("Refinement: {} residues restrained, {} held fixed{}; {} bonds, {} metal bonds ({} links made), {} "
                    "angles, {} torsions, {} chiral volumes, {} planes",
                    topo.restrained_residues, topo.fixed_residues,
                    topo.fixed_names.empty() ? std::string() : " (" + Joined(topo.fixed_names) + ")",
                    topo.bonds, topo.metal_bonds.records.size(), topo.metal_links, topo.angles, topo.torsions, topo.chirals, topo.planes);

        ctx.engine->SetRestraints(1, topo.adp);
        {
            const RestraintSet set = RestraintsNow(ctx, topo);
            ctx.engine->SetRestraints(0, set);
            RestraintSet repulsions = BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0, ctx.nthreads);
            result.geometry_start = MeasureGeometry(topo.geometry, repulsions, ctx.XyzB(), &topo.metal_bonds);
        }
        double w0 = 0, adp_w0 = 0;
        {
            const RefinementStage stage("balanced_weight");
            w0 = BalancedWeight(ctx, refinable);
            adp_w0 = BalancedADPWeight(ctx, BRestraintsNow(ctx, topo), refinable);
        }
        logger.Info("Refinement: geometry as given rmsd bonds {:.4f} A, angles {:.2f} deg, {} clashes; balanced X-ray "
                    "weights {:.3g} (positions), {:.3g} (B)", result.geometry_start.rmsd_bonds,
                    result.geometry_start.rmsd_angles, result.geometry_start.clashes, w0, adp_w0);

        const std::vector<gemmi::Position> pos0 = ctx.Positions();
        const std::vector<double> b0 = ctx.B();

        std::optional<RefinementStage> slot_stage("slots");
        for (std::thread &t : slot_builders)
            t.join();
        const size_t main_threads = ctx.nthreads;
        ctx.nthreads = slot_threads;
        std::vector<TrialSlot> trial_slots{{&ctx, &topo}};
        for (const std::unique_ptr<OwnSlot> &o : own)
            if (o)
                trial_slots.push_back({&o->ctx, &o->topo});
        slot_stage.reset();
        logger.Debug("Refinement: {} weight trials at a time", trial_slots.size());

        // A trial displaces another only where its geometry is within bounds and its R-free lower by more than the
        // tie.
        auto displaces = [](const RefinementWeightTrial &t, const RefinementWeightTrial &b) {
            return t.within_bounds && (!b.within_bounds || t.r_free < b.r_free - R_FREE_TIE);
        };
        // Of a set of trials, the one kept: they come in order of trust, the first kept unless a later one displaces
        // it. -1: none ran.
        auto choose = [&](std::vector<TrialOutcome> &outcomes, const char *stage) {
            int best = -1;
            for (size_t k = 0; k < outcomes.size(); k++) {
                const TrialOutcome &o = outcomes[k];
                if (!o.ran)
                    continue;
                logger.Info("Refinement, {}: {}", stage, o.line);
                result.trials.push_back(o.trial);
                if (best < 0 || displaces(o.trial, outcomes[best].trial))
                    best = static_cast<int>(k);
            }
            return best;
        };
        const double w_max = w0 * settings.max_weight;

        // 1. The weights from the placed model: the ladder's rungs, tighter geometry first, each with the B factors
        //    held and refined at their own balanced weight.
        std::vector<std::pair<double, double>> ladder;
        for (double rung : settings.weight_ladder)
            for (double adp : settings.adp_ladder)
                ladder.emplace_back(w0 * rung, adp_w0 * adp);
        std::vector<TrialOutcome> outcomes = RunTrials(trial_slots, ladder, refinable, pos0, b0, nullptr,
                                                       settings.macro_cycles, settings);
        int k = choose(outcomes, "ladder");
        TrialOutcome state;
        if (k >= 0)
            state = std::move(outcomes[k]);

        // 2. The trajectory: steps of a few macro-cycles from where the last one ended, each trying the weights as they
        //    are and, in turn, the positions' X-ray weight halved and doubled (halved: the geometry loosens as the
        //    model improves at one weight, until it leaves the bounds) and the B factors' divided and multiplied by
        //    four.
        //    A step is taken while it lowers R-free or R-work by more than the tie and its R-free stays within the tie
        //    of the best so far; the trajectory ends at its best step (on a deposited model, whose free set it was
        //    fitted to, R-work falls while R-free rises, and it ends where it began).
        TrialOutcome best = state;
        for (int step = 0; step < settings.max_steps && state.ran && !ctx.Cancelled(); step++) {
            const double w = state.trial.weight, adp = state.trial.adp_weight;
            std::vector<std::pair<double, double>> trials = {{w, adp}};
            if (step % 2 == 1 && adp > 0) {
                trials.emplace_back(w, adp / 4);
                trials.emplace_back(w, 4 * adp);
            } else {
                trials.emplace_back(w / 2, adp);
                if (2 * w <= w_max)
                    trials.emplace_back(2 * w, adp);
            }
            outcomes = RunTrials(trial_slots, trials, refinable, state.pos, state.b, &state.fit, settings.macro_cycles,
                                 settings);
            k = choose(outcomes, fmt::format("step {}", step + 1).c_str());
            if (k < 0)
                break;
            const RefinementWeightTrial &t = outcomes[k].trial;
            const bool gains = displaces(t, state.trial) || (t.within_bounds && t.r_work < state.trial.r_work - R_FREE_TIE);
            if (!gains || !(t.within_bounds && t.r_free < best.trial.r_free + R_FREE_TIE))
                break;
            state = std::move(outcomes[k]);
            if (displaces(state.trial, best.trial))
                best = state;
        }
        state = std::move(best);

        // 3. The final weights: a short ladder from where the trajectory ended, tighter geometry first, its choice
        //    kept where it displaces that.
        if (state.ran && !ctx.Cancelled()) {
            const double w = state.trial.weight, adp = state.trial.adp_weight;
            std::vector<std::pair<double, double>> trials;
            for (double f : settings.final_ladder)
                if (w * f <= w_max)
                    trials.emplace_back(w * f, adp);
            outcomes = RunTrials(trial_slots, trials, refinable, state.pos, state.b, &state.fit, settings.macro_cycles,
                                 settings);
            k = choose(outcomes, "final");
            if (k >= 0 && displaces(outcomes[k].trial, state.trial))
                state = std::move(outcomes[k]);
        }
        for (size_t s = 1; s < trial_slots.size(); s++)
            ctx.evaluations += trial_slots[s].ctx->evaluations;
        ctx.nthreads = main_threads;
        if (state.ran) {
            ctx.SetPositions(state.pos);
            ctx.SetB(state.b);
            PlaceRiders(ctx, topo);
            result.restrained = true;
            result.weight = state.trial.weight;
            result.adp_weight = state.trial.adp_weight;
            result.r_work = state.trial.r_work;
            result.r_free = state.trial.r_free;
            result.geometry = state.trial.geometry;
        } else
            ctx.Refit();   // no trial ran: the model as placed, its scale and solvent for the maps
    }
    result.cancelled = ctx.Cancelled();
    if (result.cancelled)
        logger.Warning("Refinement: cancelled, {} weight trials ran", result.trials.size());

    finish();
    logger.Info("Refinement: done in {:.1f} s, {} target evaluations{}", result.seconds, ctx.evaluations,
                result.restrained ? fmt::format("; R-work {:.4f}, R-free {:.4f}, weight {:.3g}", result.r_work,
                                                result.r_free, result.weight)
                                  : std::string());
    RefinementStage::Log(logger);
    return result;
}