// 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 <map>
#include <optional>
#include <cmath>
#include <random>
#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/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 "RefinementXray.h"
#include "SigmaA.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;

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;
#ifdef JFJOCH_USE_CUDA
    std::unique_ptr<ModelScaleGPUPool> scale_gpu;
#endif
    size_t nthreads = 1;
    XrayRowModel rows;
    std::vector<std::complex<float>> fcalc, fmask;
    double k_sol = 0, b_sol = 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;

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

    // 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");
            engine->Fmask(Model(), 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);
#ifdef JFJOCH_USE_CUDA
        if (scale_gpu)
            scale_gpu->Fit(scaling, SolventBMax(Model()), nthreads);
        else
#endif
            FitModelScale(scaling, {.b_hi = SolventBMax(Model())}, nthreads);
        s3.reset();
        k_sol = scaling.k_sol;
        b_sol = scaling.b_sol;
        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, nthreads);
        }
        engine->SetRowModel(rows);
        return ComputeRFactors(refl, fmodel);
    }

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

// The positions or the B factors of the refinable atoms, against w times the X-ray target plus a restraint set.
// Following Liu & Nocedal (1989) Math. Program. 45, 503-528 for the minimiser (Ceres' L-BFGS)
class AtomTarget : public ceres::FirstOrderFunction {
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) {}

    int NumParameters() const override {
        return static_cast<int>(param_atoms_.size()) * (mode_ == Mode::Xyz ? 3 : 1);
    }

    bool Evaluate(const double *p, double *cost, double *gradient) const override {
        const RefinementStage stage(mode_ == Mode::Xyz ? "eval.xyz" : "eval.b");
        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();
        ctx_.evaluations++;
        std::vector<std::array<double, 4>> gx, gr;
        double xray = 0;
        {
            const RefinementStage s1("eval.xray");
            xray = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &gx : nullptr);
        }
        if (gradient != nullptr)
            ctx_.FoldRiders(gx);
        double restraints = 0;
        if (restraint_set_ >= 0) {
            const RefinementStage s2("eval.restraints");
            restraints = ctx_.engine->EvaluateRestraints(restraint_set_, ctx_.XyzB(), gradient != nullptr ? &gr : nullptr);
        }
        *cost = weight_ * xray + restraints;
        if (gradient != nullptr)
            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] + (restraint_set_ >= 0 ? gr[i][c] : 0.0);
                else {
                    // held at a bound, B does not move further out
                    const bool held = (p[k] <= ctx_.b_floor && gx[i][3] > 0) || (p[k] >= ctx_.b_ceiling && gx[i][3] < 0);
                    gradient[k] = held ? 0.0 : weight_ * gx[i][3] + (restraint_set_ >= 0 ? gr[i][3] : 0.0);
                }
            }
        return std::isfinite(*cost);
    }

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

void Minimise(Context &ctx, AtomTarget::Mode mode, const std::vector<int> &param_atoms, double weight,
              int restraint_set, int iterations) {
    if (param_atoms.empty())
        return;
    const RefinementStage stage(mode == AtomTarget::Mode::Xyz ? "minimise.xyz" : "minimise.b");
    std::vector<double> p;
    for (int i : param_atoms) {
        const gemmi::Atom &a = *ctx.atoms[i];
        if (mode == AtomTarget::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);
    }
    ceres::GradientProblem problem(std::make_unique<AtomTarget>(ctx, mode, param_atoms, weight, restraint_set));
    ceres::GradientProblemSolver::Summary summary;
    ceres::Solve(SolverOptions(iterations), problem, p.data(), &summary);
    // the atoms where the solver ended, not where its last line search probed
    for (size_t k = 0; k < param_atoms.size(); k++) {
        gemmi::Atom &a = *ctx.atoms[param_atoms[k]];
        if (mode == AtomTarget::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();
}

// 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;
        ceres::Solve(SolverOptions(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;
        ceres::Solve(SolverOptions(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 (1992)-style gradient balancing 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 model's map coefficients and its difference map, as model validation makes them (sigma_A weighted).
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<SigmaAReflection> sa;
    for (size_t m = 0; m < ctx.refl.Size(); m++)
        sa.push_back({ctx.refl.fo[m], std::abs(fmodel[m]), ctx.refl.inv_d2[m], static_cast<int>(ctx.refl.epsilon[m]),
                      static_cast<bool>(ctx.refl.centric[m]), static_cast<bool>(ctx.refl.free[m])});
    const SigmaAResult sigma_a = EstimateSigmaA(sa, ctx.cell);
    result.mean_fom = sigma_a.mean_fom;

    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 = sigma_a.weight[m].m, d = sigma_a.weight[m].d;
        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;
    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;
        }
    }
}

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

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

    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;
    logger.Info("Refinement: {} atoms, {} reflections to {:.2f} A ({} free), on the {}; as given R-work {:.4f}, R-free {:.4f}",
                ctx.atoms.size(), ctx.refl.Size(), result.d_min, r.n_free, result.on_gpu ? "GPU" : "CPU", r.r_work,
                r.r_free);

    RefinePlacement(ctx, settings.group_b, result);
    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) {
        std::optional<RefinementStage> topo_stage("topology");
        RefinementTopology topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
        topo_stage.reset();
        // the hydrogens it added: the atoms and the engine made again for them
        ctx.atoms = AtomList(ctx.Model());
        ctx.rider_parent = topo.rider_parent;
        auto place_riders = [&]() {
            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;
                }
        };
        place_riders();
        if (static_cast<int>(ctx.atoms.size()) != result.atoms) {
            MakeEngine(ctx, result.on_gpu, settings, logger);
            result.on_gpu = ctx.engine->OnGPU();
        }
        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, {} angles, {} torsions, {} chiral "
                    "volumes, {} planes",
                    topo.restrained_residues, topo.fixed_residues,
                    topo.fixed_names.empty() ? std::string() : " (" + Joined(topo.fixed_names) + ")",
                    topo.bonds, topo.angles, topo.torsions, topo.chirals, topo.planes);

        auto restraints_now = [&]() {
            const RefinementStage stage("repulsions");
            RestraintSet set = topo.geometry;
            set.Append(BuildRepulsions(topo, ctx.Model(), cell, sg, REPULSION_MARGIN));
            set.Finalize(ctx.atoms.size());
            return set;
        };
        ctx.engine->SetRestraints(1, topo.adp);
        {
            const RestraintSet set = restraints_now();
            ctx.engine->SetRestraints(0, set);
            RestraintSet repulsions = BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0);
            result.geometry_start = MeasureGeometry(topo.geometry, repulsions, ctx.XyzB());
        }
        double w0 = 0;
        {
            const RefinementStage stage("balanced_weight");
            w0 = BalancedWeight(ctx, refinable);
        }
        logger.Info("Refinement: geometry as given rmsd bonds {:.4f} A, angles {:.2f} deg, {} clashes; balanced X-ray "
                    "weight {:.3g}", result.geometry_start.rmsd_bonds, result.geometry_start.rmsd_angles,
                    result.geometry_start.clashes, w0);

        const std::vector<gemmi::Position> pos0 = ctx.Positions();
        const std::vector<double> b0 = ctx.B();
        int best = -1;
        std::vector<gemmi::Position> best_pos;
        std::vector<double> best_b;
        // One full refinement from the start: macro-cycles of the positions at X-ray weight w and, where adp_w is
        // not 0, the B factors at X-ray weight adp_w; kept where it beats the best so far. The rungs come in order
        // of trust - tighter geometry, then fewer parameters - and a later one displaces an earlier only by more
        // than the tie in R-free.
        auto trial = [&](double w, double adp_w) {
            const RefinementStage stage(adp_w > 0 ? "rung.b" : "rung.xyz");
            ctx.SetPositions(pos0);
            ctx.SetB(b0);
            for (int cycle = 0; cycle < settings.macro_cycles; cycle++) {
                place_riders();
                ctx.Refit();
                {
                    const RestraintSet set = restraints_now();
                    const RefinementStage s1("set_restraints");
                    ctx.engine->SetRestraints(0, set);
                }
                Minimise(ctx, AtomTarget::Mode::Xyz, refinable, w, 0, settings.xyz_iterations);
                if (adp_w > 0)
                    Minimise(ctx, AtomTarget::Mode::B, refinable, adp_w, 1, settings.b_iterations);
            }
            place_riders();
            RefinementWeightTrial t;
            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(), cell, sg, 0.0), ctx.XyzB());
            }
            t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES;
            logger.Info("Refinement, 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.within_bounds ? "" : " (geometry out of bounds)");
            result.trials.push_back(t);
            const int k = static_cast<int>(result.trials.size()) - 1;
            const bool better = best < 0 ? true
                                : t.within_bounds && (!result.trials[best].within_bounds ||
                                                      t.r_free < result.trials[best].r_free - R_FREE_TIE);
            if (better) {
                best = k;
                best_pos = ctx.Positions();
                best_b = ctx.B();
            }
        };
        // the positions' weight, with the B factors as they are
        for (double ladder : settings.weight_ladder)
            trial(w0 * ladder, 0.0);
        // then the B factors refined too, at that weight, against their own weight
        const double w = result.trials[best].weight;
        for (double adp : settings.adp_ladder)
            trial(w, w * adp);
        ctx.SetPositions(best_pos);
        ctx.SetB(best_b);
        place_riders();
        const RefinementWeightTrial &chosen = result.trials[best];
        result.restrained = true;
        result.weight = chosen.weight;
        result.adp_weight = chosen.adp_weight;
        result.r_work = chosen.r_work;
        result.r_free = chosen.r_free;
        result.geometry = chosen.geometry;
    }

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