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

#include <catch2/catch_all.hpp>

#include <algorithm>
#include <atomic>
#include <chrono>
#include <cmath>
#include <map>
#include <random>
#include <set>
#include <vector>

#include <gemmi/mmread_gz.hpp>
#include <gemmi/neighbor.hpp>
#include <gemmi/modify.hpp>
#include <gemmi/symmetry.hpp>

#include "../common/CUDAWrapper.h"
#include "../common/Logger.h"
#include "../common/UnitCell.h"
#include "../image_analysis/scale_merge/RfreeFlags.h"
#include "../image_analysis/structure_refinement/ModelGrid.h"
#include "../image_analysis/structure_refinement/MonomerLibrary.h"
#include "../image_analysis/structure_refinement/Refinement.h"
#include "../image_analysis/structure_refinement/RefinementADP.h"
#include "../image_analysis/structure_refinement/RefinementMath.h"
#include "../image_analysis/structure_refinement/RefinementMinimiser.h"
#include "../image_analysis/structure_refinement/RefinementNcs.h"
#include "../image_analysis/structure_refinement/RefinementRestraints.h"
#include "../image_analysis/structure_refinement/RefinementXray.h"
#ifdef JFJOCH_USE_CUDA
#include "../image_analysis/structure_refinement/RefinementGPU.h"
#endif

namespace {

// Copies of residue `code` at the embedded dictionary's ideal coordinates (hydrogens left out), each turned and
// moved to a place of its own in the cell.
gemmi::Structure residue_crystal(const std::vector<std::string> &codes, const gemmi::UnitCell &cell, const char *group,
                                 unsigned seed) {
    std::mt19937 rng(seed);
    std::uniform_real_distribution<double> frac(0.0, 1.0), angle(-3.1, 3.1);
    gemmi::Structure st;
    st.cell = cell;
    st.spacegroup_hm = group;
    st.setup_cell_images();
    gemmi::Model model(1);
    gemmi::Chain chain("A");
    int seq = 1;
    for (const std::string &code : codes) {
        const gemmi::ChemComp &cc = EmbeddedMonomerLibrary().monomers.at(code);
        gemmi::Residue res;
        res.name = code;
        res.seqid = gemmi::SeqId(seq++, ' ');
        res.het_flag = 'H';
        res.entity_type = gemmi::EntityType::NonPolymer;
        const gemmi::Position centre = cell.orthogonalize(gemmi::Fractional(frac(rng), frac(rng), frac(rng)));
        const double ax = angle(rng), ay = angle(rng), az = angle(rng);
        const gemmi::Mat33 rot = gemmi::Mat33(1, 0, 0, 0, std::cos(ax), -std::sin(ax), 0, std::sin(ax), std::cos(ax))
                                     .multiply(gemmi::Mat33(std::cos(ay), 0, std::sin(ay), 0, 1, 0, -std::sin(ay), 0, std::cos(ay)))
                                     .multiply(gemmi::Mat33(std::cos(az), -std::sin(az), 0, std::sin(az), std::cos(az), 0, 0, 0, 1));
        for (const gemmi::ChemComp::Atom &ca : cc.atoms) {
            if (ca.is_hydrogen())
                continue;
            gemmi::Atom a;
            a.name = ca.id;
            a.element = ca.el;
            a.pos = gemmi::Position(rot.multiply(ca.xyz)) + centre;
            a.b_iso = 20.0f;
            a.occ = 1.0f;
            res.atoms.push_back(a);
        }
        chain.residues.push_back(res);
    }
    model.chains.push_back(chain);
    st.models.push_back(model);
    return st;
}

// Atoms of a few elements scattered over the cell.
gemmi::Structure random_crystal(const gemmi::UnitCell &cell, const char *group, int n, unsigned seed) {
    std::mt19937 rng(seed);
    std::uniform_real_distribution<double> frac(0.0, 1.0), b(10.0, 40.0);
    const gemmi::El elements[] = {gemmi::El::C, gemmi::El::N, gemmi::El::O, gemmi::El::S};
    gemmi::Structure st;
    st.cell = cell;
    st.spacegroup_hm = group;
    st.setup_cell_images();
    gemmi::Model model(1);
    gemmi::Chain chain("A");
    gemmi::Residue res;
    res.name = "UNK";
    res.seqid = gemmi::SeqId(1, ' ');
    for (int i = 0; i < n; i++) {
        gemmi::Atom a;
        a.name = "X";
        a.element = elements[i % 4];
        a.pos = cell.orthogonalize(gemmi::Fractional(frac(rng), frac(rng), frac(rng)));
        a.b_iso = static_cast<float>(b(rng));
        a.occ = 1.0f;
        res.atoms.push_back(a);
    }
    chain.residues.push_back(res);
    model.chains.push_back(chain);
    st.models.push_back(model);
    return st;
}

// Every reflection of the group's ASU to d_min, amplitude 1, one in twenty free.
std::vector<MergedReflection> asu_reflections(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) {
    std::vector<MergedReflection> out;
    const gemmi::ReciprocalAsu asu(&sg);
    const gemmi::GroupOps gops = sg.operations();
    const gemmi::Miller lim = cell.get_hkl_limits(d_min);
    int i = 0;
    for (int h = -lim[0]; h <= lim[0]; h++)
        for (int k = -lim[1]; k <= lim[1]; k++)
            for (int l = -lim[2]; l <= lim[2]; l++) {
                const gemmi::Miller m{{h, k, l}};
                if ((h == 0 && k == 0 && l == 0) || !asu.is_in(m) || gops.is_systematically_absent(m))
                    continue;
                const double d = cell.calculate_d(m);
                if (d < d_min)
                    continue;
                MergedReflection r;
                r.h = h;
                r.k = k;
                r.l = l;
                r.d = static_cast<float>(d);
                r.F = 1.0f;
                r.sigmaF = 0.05f;
                r.rfree_flag = (i++ % 20) == 0;
                out.push_back(r);
            }
    return out;
}

// The observations set to |F| of `truth`, and a row model with no scale, no solvent and the likelihood's
// parameters estimated against `model`.
void set_observations(RefinementEngine &engine, RefinementReflections &refl, const gemmi::Model &truth,
                      const gemmi::Model &model) {
    std::vector<std::complex<float>> f;
    engine.Fcalc(truth, f);
    for (size_t m = 0; m < refl.Size(); m++)
        refl.fo[m] = std::abs(f[m]);
    engine.Fcalc(model, f);
    XrayRowModel rows;
    rows.k_total.assign(refl.Size(), 1.0f);
    rows.fmask.assign(refl.Size(), {0, 0});
    EstimateLikelihoodParameters(refl, f, rows);
    engine.SetRowModel(rows);
}

gemmi::Model shaken(const gemmi::Model &model, double sigma, double b_sigma, unsigned seed) {
    std::mt19937 rng(seed);
    std::normal_distribution<double> d(0, sigma), db(0, b_sigma);
    gemmi::Model out = model;
    for (gemmi::Chain &ch : out.chains)
        for (gemmi::Residue &r : ch.residues)
            for (gemmi::Atom &a : r.atoms) {
                a.pos += gemmi::Position(d(rng), d(rng), d(rng));
                a.b_iso = static_cast<float>(std::max(5.0, a.b_iso + db(rng)));
            }
    return out;
}

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

// The engine's gradient of a few atoms against central differences of its own target.
void check_xray_gradient(RefinementEngine &engine, gemmi::Model model, double tolerance) {
    std::vector<std::array<double, 4>> grad;
    engine.Evaluate(model, &grad);
    std::vector<gemmi::Atom *> atoms = atoms_of(model);
    const double hx = 2e-3, hb = 2e-2;
    double max_g = 0;
    for (const auto &g : grad)
        for (int c = 0; c < 4; c++)
            max_g = std::max(max_g, std::fabs(g[c]));
    for (size_t i = 0; i < atoms.size(); i += atoms.size() / 7) {
        for (int c = 0; c < 4; c++) {
            const gemmi::Position p0 = atoms[i]->pos;
            const float b0 = atoms[i]->b_iso;
            const double h = c < 3 ? hx : hb;
            auto set = [&](double s) {
                atoms[i]->pos = p0;
                atoms[i]->b_iso = b0;
                if (c < 3)
                    atoms[i]->pos.at(c) += s;
                else
                    atoms[i]->b_iso = static_cast<float>(b0 + s);
            };
            set(h);
            const double fp = engine.Evaluate(model, nullptr);
            set(-h);
            const double fm = engine.Evaluate(model, nullptr);
            set(0);
            const double fd = (fp - fm) / (2 * h);
            CAPTURE(i, c, fd, grad[i][c], max_g);
            CHECK(std::fabs(fd - grad[i][c]) < tolerance * std::max(std::fabs(fd), 0.05 * max_g));
        }
    }
}

}  // namespace

TEST_CASE("RefinementLikelihood_derivative", "[Refinement]") {
    for (bool centric : {false, true})
        for (double fc : {0.3, 2.0, 9.0}) {
            double dfc = 0, dummy = 0;
            const double f0 = refine_math::AmplitudeNegLogLikelihood<double>(4.0, fc, 0.8, 2.5, centric, dfc);
            const double h = 1e-5;
            const double fp = refine_math::AmplitudeNegLogLikelihood<double>(4.0, fc + h, 0.8, 2.5, centric, dummy);
            const double fm = refine_math::AmplitudeNegLogLikelihood<double>(4.0, fc - h, 0.8, 2.5, centric, dummy);
            CAPTURE(centric, fc, f0);
            CHECK(dfc == Catch::Approx((fp - fm) / (2 * h)).epsilon(1e-5));
        }
    // the Bessel approximation on both sides of its switch
    double l0, r0, l1, r1;
    refine_math::BesselI0I1(3.75 * (1 - 1e-12), l0, r0);
    refine_math::BesselI0I1(3.75 * (1 + 1e-12), l1, r1);
    CHECK(l0 == Catch::Approx(l1).epsilon(1e-5));
    CHECK(r0 == Catch::Approx(r1).epsilon(1e-5));
}

TEST_CASE("RefinementMonomerLibrary_embedded", "[Refinement]") {
    const gemmi::MonLib &monlib = EmbeddedMonomerLibrary();
    for (const char *code : {"ALA", "PRO", "MSE", "DA", "U", "ATP", "GOL", "SO4", "HOH", "ZN"})
        CHECK(InEmbeddedMonomerLibrary(code));
    CHECK_FALSE(InEmbeddedMonomerLibrary("XYZ"));
    for (const char *link : {"TRANS", "CIS", "PTRANS", "p", "SS"})
        CHECK(monlib.links.count(link) == 1);
    CHECK(monlib.modifications.count("NH3") == 1);
    CHECK(monlib.ener_lib.atoms.count("CH1") == 1);
}

TEST_CASE("RefinementRestraints_ideal_residues", "[Refinement]") {
    const gemmi::UnitCell cell(60, 65, 70, 90, 90, 90);
    gemmi::Structure st = residue_crystal({"PHE", "DA", "ATP", "HOH", "ZN"}, cell, "P 21 21 21", 3);
    const RefinementTopology topo = BuildRefinementTopology(st);
    CHECK(topo.restrained_residues == 5);
    CHECK(topo.fixed_residues == 0);
    CHECK(topo.bonds > 60);
    CHECK(topo.planes >= 3);
    CHECK(topo.chirals >= 4);
    std::vector<std::array<double, 4>> xyzb;
    for (gemmi::Atom *a : atoms_of(st.models[0]))
        xyzb.push_back({a->pos.x, a->pos.y, a->pos.z, a->b_iso});
    const GeometryStatistics g = MeasureGeometry(topo.geometry, {}, xyzb);
    CHECK(g.rmsd_bonds < 0.02);
    CHECK(g.rmsd_angles < 2.5);
    CHECK(g.rmsd_chirals < 0.2);
}

// A zinc bound by four cysteines (the dictionary's own ZN-CYS link) and a magnesium by six waters (links made from
// the library's metal-site statistics), both from LINK records: the bonds come out at the expected distances,
// and a metal set off its site returns there under them alone.
TEST_CASE("RefinementRestraints_metal_links", "[Refinement]") {
    const gemmi::UnitCell cell(40, 40, 40, 90, 90, 90);
    gemmi::Structure st = residue_crystal({"CYS", "CYS", "CYS", "CYS", "ZN", "HOH", "HOH", "HOH", "HOH", "HOH", "HOH",
                                           "MG"}, cell, "P 1", 11);
    double mg_o = 0, mg_sigma = 0;
    REQUIRE(EmbeddedMetalDistance(gemmi::El::Mg, gemmi::El::O, 6, mg_o, mg_sigma));
    CHECK(mg_o == Catch::Approx(2.07).margin(0.02));   // Mg-O of an octahedral site
    constexpr double ZN_S = 2.34;                      // the ZN-CYS link
    std::vector<gemmi::Residue> &res = st.models[0].chains[0].residues;
    // the donors on a tetrahedron and an octahedron around their sites, each residue moved as a whole
    const gemmi::Position zn_site(10, 10, 10), mg_site(25, 25, 25);
    const double t = 1.0 / std::sqrt(3.0);
    const gemmi::Position tetra[4] = {{t, t, t}, {t, -t, -t}, {-t, t, -t}, {-t, -t, t}};
    const gemmi::Position octa[6] = {{1, 0, 0}, {-1, 0, 0}, {0, 1, 0}, {0, -1, 0}, {0, 0, 1}, {0, 0, -1}};
    auto donor_to = [](gemmi::Residue &r, const std::string &atom, const gemmi::Position &target) {
        const gemmi::Position shift = target - r.find_atom(atom, '*')->pos;
        for (gemmi::Atom &a : r.atoms)
            a.pos += shift;
    };
    for (int k = 0; k < 4; k++)
        donor_to(res[k], "SG", zn_site + tetra[k] * ZN_S);
    for (int k = 0; k < 6; k++)
        donor_to(res[5 + k], "O", mg_site + octa[k] * mg_o);
    // the metals off their sites by 0.3 A
    res[4].atoms[0].pos = zn_site + gemmi::Position(0.2, -0.15, 0.17);
    res[11].atoms[0].pos = mg_site + gemmi::Position(-0.1, 0.25, 0.12);
    auto link = [&](int metal, int donor, const char *atom) {
        gemmi::Connection c;
        c.type = gemmi::Connection::MetalC;
        c.asu = gemmi::Asu::Same;
        c.partner1 = gemmi::AtomAddress("A", res[donor].seqid, res[donor].name, atom);
        c.partner2 = gemmi::AtomAddress("A", res[metal].seqid, res[metal].name, res[metal].atoms[0].name);
        st.connections.push_back(c);
    };
    for (int k = 0; k < 4; k++)
        link(4, k, "SG");
    for (int k = 0; k < 6; k++)
        link(11, 5 + k, "O");

    const RefinementTopology topo = BuildRefinementTopology(st);
    CHECK(topo.metal_links == 6);   // the Mg-O links; Zn-S is the dictionary's
    std::vector<gemmi::Atom *> atoms = atoms_of(st.models[0]);
    const int zn = static_cast<int>(std::find(atoms.begin(), atoms.end(), &res[4].atoms[0]) - atoms.begin());
    const int mg = static_cast<int>(std::find(atoms.begin(), atoms.end(), &res[11].atoms[0]) - atoms.begin());
    // the metals' bonds kept apart from the covalent ones, none left among them
    for (const RestraintRecord &r : topo.geometry.records)
        CHECK_FALSE((r.kind == static_cast<int>(RestraintKind::Bond) &&
                     (r.atom[0] == zn || r.atom[1] == zn || r.atom[0] == mg || r.atom[1] == mg)));
    const RestraintSet &metal_bonds = topo.metal_bonds;
    REQUIRE(metal_bonds.records.size() == 10);
    for (const RestraintRecord &r : metal_bonds.records) {
        const bool is_zn = r.atom[0] == zn || r.atom[1] == zn;
        CHECK((is_zn || r.atom[0] == mg || r.atom[1] == mg));
        CHECK(r.ideal == Catch::Approx(is_zn ? ZN_S : mg_o).margin(1e-4));
    }

    // the two metals moved down the gradient of their bonds, everything else held
    std::vector<std::array<double, 4>> xyzb;
    for (gemmi::Atom *a : atoms)
        xyzb.push_back({a->pos.x, a->pos.y, a->pos.z, a->b_iso});
    double e = EvaluateRestraints(metal_bonds, xyzb, nullptr), step = 1e-4;
    for (int it = 0; it < 2000 && step > 1e-12; it++) {
        std::vector<std::array<double, 4>> grad(xyzb.size(), {0, 0, 0, 0}), trial = xyzb;
        EvaluateRestraints(metal_bonds, xyzb, &grad);
        for (int a : {zn, mg})
            for (int k = 0; k < 3; k++)
                trial[a][k] -= step * grad[a][k];
        const double e_trial = EvaluateRestraints(metal_bonds, trial, nullptr);
        if (e_trial < e) {
            xyzb = trial;
            e = e_trial;
            step *= 1.5;
        } else
            step *= 0.5;
    }
    for (const RestraintRecord &r : metal_bonds.records)
        CHECK(RestraintValue(metal_bonds, r, xyzb) == Catch::Approx(r.ideal).margin(0.005));
}

// A magnesium with six waters at its inner-sphere distance and a seventh at 2.9 A, all seven from LINK records: the
// six become bonds, the seventh - a record at an outer-sphere distance - stays a contact, and a second topology of
// the same model (a second round of refinement) finds the same, the records keeping the names they were read with.
TEST_CASE("RefinementRestraints_metal_links_outer_sphere", "[Refinement]") {
    const gemmi::UnitCell cell(40, 40, 40, 90, 90, 90);
    gemmi::Structure st = residue_crystal({"HOH", "HOH", "HOH", "HOH", "HOH", "HOH", "HOH", "MG"}, cell, "P 1", 7);
    std::vector<gemmi::Residue> &res = st.models[0].chains[0].residues;
    const gemmi::Position site(20, 20, 20);
    const gemmi::Position octa[6] = {{1, 0, 0}, {-1, 0, 0}, {0, 1, 0}, {0, -1, 0}, {0, 0, 1}, {0, 0, -1}};
    for (int k = 0; k < 6; k++)
        res[k].atoms[0].pos = site + octa[k] * 2.07;
    res[6].atoms[0].pos = site + gemmi::Position(1, 1, 1) * (2.9 / std::sqrt(3.0));
    res[7].atoms[0].pos = site;
    for (int k = 0; k < 7; k++) {
        gemmi::Connection c;
        c.type = gemmi::Connection::MetalC;
        c.asu = gemmi::Asu::Same;
        c.partner1 = gemmi::AtomAddress("A", res[k].seqid, "HOH", "O");
        c.partner2 = gemmi::AtomAddress("A", res[7].seqid, "MG", res[7].atoms[0].name);
        st.connections.push_back(c);
    }
    std::vector<gemmi::Atom *> atoms = atoms_of(st.models[0]);
    auto index = [&](const gemmi::Atom &a) {
        return static_cast<int>(std::find(atoms.begin(), atoms.end(), &a) - atoms.begin());
    };
    const int mg = index(res[7].atoms[0]), outer = index(res[6].atoms[0]);
    for (int round = 0; round < 2; round++) {
        const RefinementTopology topo = BuildRefinementTopology(st);
        CHECK(topo.metal_links == 6);
        CHECK(topo.metal_links_refused == 1);
        REQUIRE(topo.metal_bonds.records.size() == 6);
        for (const RestraintRecord &r : topo.metal_bonds.records)
            CHECK((r.atom[0] != outer && r.atom[1] != outer));
        for (const gemmi::Connection &c : st.connections)
            CHECK(c.link_id.empty());
        // the outer water is a contact of the metal's (listed with a margin that reaches it)
        const RestraintSet contacts = BuildRepulsions(topo, st.models[0], st.cell, *st.find_spacegroup(), 1.5);
        int found = 0;
        for (const RestraintRecord &r : contacts.records)
            found += (r.atom[0] == mg && r.atom[1] == outer) || (r.atom[0] == outer && r.atom[1] == mg);
        CHECK(found == 1);
    }
}

// Deposited models cut down to a ligand and the residues within 4.2 A of it (test_data/ligands, from the PDB entries
// their names start with; the tRNA as its modified nucleotides and their neighbours in the chain): every ligand is
// restrained, its own restraints included, the deposited geometry agrees with the dictionary, and the links the
// ligands need are made - the RNA backbone to and from each modified nucleotide, the iron of a heme to its histidine
// (and a heme C's to its methionine) from the LINK records, and a heme C's two thioethers to its cysteines.
TEST_CASE("RefinementRestraints_deposited_ligands", "[Refinement]") {
    const std::vector<std::string> modified_nucleotides = {"PSU", "5MC", "1MA", "OMG", "OMC", "H2U", "M2G", "7MG",
                                                           "2MG", "5MU"};
    const std::vector<std::pair<std::string, std::vector<std::string>>> cases = {
        {"1ehz_psu", modified_nucleotides}, {"8bkh_hem", {"HEM"}}, {"6k9i_hec", {"HEC"}}, {"2w3v_ndp", {"NDP"}},
        {"5vml_nap", {"NAP"}}, {"6f3p_nad", {"NAD"}}, {"8con_nai", {"NAI"}}, {"3tk0_fad", {"FAD"}},
        {"7tcd_fmn", {"FMN"}}, {"6w4h_sam", {"SAM"}}, {"5epe_sah", {"SAH"}}};
    for (const auto &[name, ligands] : cases) {
        DYNAMIC_SECTION(name) {
            gemmi::Structure st = gemmi::read_structure_gz("../../tests/test_data/ligands/" + name + ".pdb.gz");
            const RefinementTopology topo = BuildRefinementTopology(st);
            std::vector<gemmi::Atom *> atoms = atoms_of(st.models[0]);
            std::vector<const gemmi::Residue *> residue_of;
            for (const gemmi::Chain &ch : st.models[0].chains)
                for (const gemmi::Residue &r : ch.residues)
                    for (size_t k = 0; k < r.atoms.size(); k++)
                        residue_of.push_back(&r);
            auto index_of = [&](const gemmi::Residue &r, const char *atom) {
                const gemmi::Atom *a = r.find_atom(atom, '*');
                return a == nullptr ? -1 : static_cast<int>(std::find(atoms.begin(), atoms.end(), a) - atoms.begin());
            };
            auto bonded = [&](const RestraintSet &set, int i, int j) {
                for (const RestraintRecord &r : set.records)
                    if (r.kind == static_cast<int>(RestraintKind::Bond) &&
                        ((r.atom[0] == i && r.atom[1] == j) || (r.atom[0] == j && r.atom[1] == i)))
                        return true;
                return false;
            };

            std::map<std::string, int> restraints, found;
            for (const RestraintRecord &r : topo.geometry.records)
                if (r.kind != static_cast<int>(RestraintKind::Plane))
                    restraints[residue_of[r.atom[0]]->name]++;
            for (size_t i = 0; i < atoms.size(); i++)
                if (std::find(ligands.begin(), ligands.end(), residue_of[i]->name) != ligands.end()) {
                    found[residue_of[i]->name]++;
                    CHECK(topo.refinable[i]);
                }
            for (const std::string &code : ligands) {
                INFO(code);
                CHECK(found[code] > 0);
                CHECK(restraints[code] > 0);
                CHECK(std::find(topo.fixed_names.begin(), topo.fixed_names.end(), code) == topo.fixed_names.end());
            }

            std::vector<std::array<double, 4>> xyzb;
            for (gemmi::Atom *a : atoms)
                xyzb.push_back({a->pos.x, a->pos.y, a->pos.z, a->b_iso});
            const GeometryStatistics g = MeasureGeometry(topo.geometry, {}, xyzb, &topo.metal_bonds);
            CHECK(g.rmsd_bonds < 0.03);
            CHECK(g.rmsd_angles < 3.0);
            CHECK(g.rmsd_metal < 0.1);

            if (name == "1ehz_psu") {
                // each modified nucleotide bonded to the residue before it in the chain, where that one is there
                int links = 0;
                for (const gemmi::Residue &r : st.models[0].chains[0].residues) {
                    const gemmi::Residue *prev = &r == &st.models[0].chains[0].residues.front() ? nullptr : &r - 1;
                    if (prev == nullptr || prev->seqid.num != r.seqid.num - 1 ||
                        std::find(ligands.begin(), ligands.end(), r.name) == ligands.end())
                        continue;
                    INFO(r.name << " " << *r.seqid.num);
                    CHECK(bonded(topo.geometry, index_of(*prev, "O3'"), index_of(r, "P")));
                    links++;
                }
                CHECK(links >= 10);
            }
            if (name == "8bkh_hem" || name == "6k9i_hec") {
                const gemmi::Residue *heme = nullptr, *his = nullptr;
                for (const gemmi::Residue &r : st.models[0].chains[0].residues) {
                    if (r.name == ligands[0])
                        heme = &r;
                    if (r.name == "HIS" && r.seqid.num == (name == "8bkh_hem" ? 93 : 18))
                        his = &r;
                }
                REQUIRE(heme != nullptr);
                REQUIRE(his != nullptr);
                const int fe = index_of(*heme, "FE");
                CHECK(topo.metal[fe]);
                CHECK(topo.metal_links == (name == "8bkh_hem" ? 1 : 2));   // the histidine; and the methionine
                CHECK(bonded(topo.metal_bonds, fe, index_of(*his, "NE2")));
                CHECK(bonded(topo.metal_bonds, fe, index_of(*heme, "NA")));   // the porphyrin's own, among the metal's
                if (name == "6k9i_hec")
                    for (auto [cys, carbon] : {std::pair{14, "CAB"}, std::pair{17, "CAC"}})
                        for (const gemmi::Residue &r : st.models[0].chains[0].residues)
                            if (r.name == "CYS" && r.seqid.num == cys)
                                CHECK(bonded(topo.geometry, index_of(r, "SG"), index_of(*heme, carbon)));
            }
        }
    }
}

TEST_CASE("RefinementRestraints_gradient", "[Refinement]") {
    const gemmi::UnitCell cell(14, 15, 16, 90, 90, 90);
    gemmi::Structure st = residue_crystal({"TYR", "DG"}, cell, "P 21 21 21", 5);
    st.models[0] = shaken(st.models[0], 0.15, 3.0, 2);
    const RefinementTopology topo = BuildRefinementTopology(st);
    // the geometry, the B similarity and the repulsions (symmetry mates included, at a contact distance that
    // makes every pair within reach count)
    RestraintSet set = topo.geometry;
    set.Append(topo.adp);
    RefinementTopology wide = topo;
    for (float &r : wide.vdw_radius)
        r = 2.5f;
    std::fill(wide.hbond.begin(), wide.hbond.end(), 0);
    const RestraintSet rep = BuildRepulsions(wide, st.models[0], cell, *st.find_spacegroup(), 0.0);
    REQUIRE(rep.records.size() > 10);
    bool image = false;
    for (const RestraintRecord &r : rep.records)
        image = image || r.rot[0] != 1.0f || r.tran[0] != 0.0f || r.tran[1] != 0.0f || r.tran[2] != 0.0f;
    CAPTURE(rep.records.size());
    CHECK(image);
    set.Append(rep);
    set.Finalize(atoms_of(st.models[0]).size());

    std::vector<std::array<double, 4>> xyzb;
    for (gemmi::Atom *a : atoms_of(st.models[0]))
        xyzb.push_back({a->pos.x, a->pos.y, a->pos.z, a->b_iso});
    // each kind of restraint on its own, against central differences
    for (int kind = 0; kind <= static_cast<int>(RestraintKind::BSimilarity); kind++) {
        RestraintSet one;
        for (const RestraintRecord &r : set.records)
            if (r.kind == kind) {
                RestraintRecord copy = r;
                if (kind == static_cast<int>(RestraintKind::Plane)) {
                    copy.atom[0] = static_cast<int>(one.plane_atoms.size());
                    for (int k = 0; k < r.atom[1]; k++)
                        one.plane_atoms.push_back(set.plane_atoms[r.atom[0] + k]);
                }
                one.Add(copy, kind == static_cast<int>(RestraintKind::Plane) ? r.atom[1] : 4);
            }
        CAPTURE(kind, one.records.size());
        std::vector<std::array<double, 4>> grad(xyzb.size(), {0, 0, 0, 0});
        EvaluateRestraints(one, xyzb, &grad);
        const double h = 1e-6;
        for (size_t i = 0; i < xyzb.size(); i++)
            for (int c = 0; c < 4; c++) {
                auto x = xyzb;
                x[i][c] += h;
                const double fp = EvaluateRestraints(one, x, nullptr);
                x[i][c] -= 2 * h;
                const double fm = EvaluateRestraints(one, x, nullptr);
                CAPTURE(i, c);
                CHECK(grad[i][c] == Catch::Approx((fp - fm) / (2 * h)).epsilon(1e-4).margin(1e-3));
            }
    }

#ifdef JFJOCH_USE_CUDA
    if (get_gpu_count() > 0) {
        // the GPU's float evaluation of the same set
        // (the engine is made for a cell wide enough for its gather; the restraints do not read the cell)
        const gemmi::UnitCell wide_cell(40, 40, 40, 90, 90, 90);
        std::vector<MergedReflection> merged = asu_reflections(wide_cell, *st.find_spacegroup(), 3.0);
        RefinementReflections refl = PrepareRefinementReflections(merged, wide_cell, *st.find_spacegroup(), 3.0);
        auto gpu = MakeRefinementEngineGPU(0, wide_cell, *st.find_spacegroup(), 3.0, refl, st.models[0], 100.0, SIZE_MAX);
        REQUIRE(gpu);
        gpu->SetRestraints(0, set);
        std::vector<std::array<double, 4>> ggpu;
        const double egpu = gpu->EvaluateRestraints(0, xyzb, &ggpu);
        std::vector<std::array<double, 4>> gcpu(xyzb.size(), {0, 0, 0, 0});
        const double ecpu = EvaluateRestraints(set, xyzb, &gcpu);
        CHECK(egpu == Catch::Approx(ecpu).epsilon(1e-4));
        for (size_t i = 0; i < xyzb.size(); i++)
            for (int c = 0; c < 4; c++)
                CHECK(ggpu[i][c] == Catch::Approx(gcpu[i][c]).epsilon(2e-3).margin(2e-2));
    }
#endif
}

// Contacts with other copies of the molecule: in P1 two waters 2.4 A apart across the a face (17.6 A apart within
// the cell) repel each other once, and a water 1.4 A from a 3-fold axis, 2.4 A from its images x -> T x and
// x -> T^-1 x (one contact per asymmetric unit), repels its image once.
TEST_CASE("RefinementRestraints_lattice_and_symmetry_contacts", "[Refinement]") {
    auto oxygens = [](gemmi::Structure &st, const std::vector<gemmi::Fractional> &where) {
        std::vector<gemmi::Atom *> atoms = atoms_of(st.models[0]);
        REQUIRE(atoms.size() == where.size());
        for (size_t i = 0; i < where.size(); i++)
            atoms[i]->pos = st.cell.orthogonalize(where[i]);
    };
    {
        const gemmi::UnitCell cell(20, 21, 22, 90, 90, 90);
        gemmi::Structure st = residue_crystal({"HOH", "HOH"}, cell, "P 1", 1);
        oxygens(st, {gemmi::Fractional(0.01, 0.5, 0.5), gemmi::Fractional(0.89, 0.5, 0.5)});
        const RefinementTopology topo = BuildRefinementTopology(st);
        const RestraintSet rep = BuildRepulsions(topo, st.models[0], cell, *st.find_spacegroup(), 0.0);
        REQUIRE(rep.records.size() == 1);
        CHECK(rep.records[0].tran[0] == Catch::Approx(-20.0f));
        CHECK(BuildBContacts(topo, st.models[0], cell, *st.find_spacegroup()).records.size() == 1);
    }
    {
        const gemmi::UnitCell cell(20, 20, 22, 90, 90, 120);
        gemmi::Structure st = residue_crystal({"HOH"}, cell, "P 3", 1);
        oxygens(st, {cell.fractionalize(gemmi::Position(1.4, 0.0, 11.0))});
        const RefinementTopology topo = BuildRefinementTopology(st);
        const RestraintSet rep = BuildRepulsions(topo, st.models[0], cell, *st.find_spacegroup(), 0.0);
        CHECK(rep.records.size() == 1);
    }
}

// A complete model whose every atom is off by a Gaussian error of sigma_x per axis: D(s) = <cos(2 pi s.dx)> =
// exp(-2 pi^2 sigma_x^2 s^2) (Luzzati). With a negligible measurement error D comes out at that; with an error of
// half the rms amplitude of the outermost shell it falls below it there, where the error is a large part of the
// spread and is counted twice, and stays where the amplitudes are strong.
TEST_CASE("RefinementLikelihood_D_coordinate_error", "[Refinement]") {
    const char *group = "P 1";
    const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group);
    const gemmi::UnitCell cell(40, 45, 50, 90, 90, 90);
    const double d_min = 1.5, sigma_x = 0.25;
    gemmi::Structure truth = random_crystal(cell, group, 1500, 3);
    gemmi::Model model = truth.models[0];
    std::mt19937 rng(4);
    std::normal_distribution<double> dx(0, sigma_x);
    for (gemmi::Atom &a : model.chains[0].residues[0].atoms)
        a.pos += gemmi::Position(dx(rng), dx(rng), dx(rng));
    RefinementReflections refl = PrepareRefinementReflections(asu_reflections(cell, sg, d_min), cell, sg, d_min);
    auto cpu = MakeRefinementEngineCPU(cell, sg, d_min, refl, 4);
    cpu->SetBlur(RefinementBlur(cell, sg, d_min, 5.0));
    std::vector<std::complex<float>> f;
    cpu->Fcalc(truth.models[0], f);
    for (size_t m = 0; m < refl.Size(); m++)
        refl.fo[m] = std::abs(f[m]);
    cpu->Fcalc(model, f);

    // mean D and mean Luzzati D in five shells of 1/d^2
    constexpr int SHELLS = 5;
    const double s2_max = 1.0 / (d_min * d_min);
    auto shells = [&](const XrayRowModel &rows, std::vector<double> &d, std::vector<double> &luzzati) {
        std::vector<int> n(SHELLS, 0);
        d.assign(SHELLS, 0);
        luzzati.assign(SHELLS, 0);
        for (size_t m = 0; m < refl.Size(); m++) {
            const int k = std::min(SHELLS - 1, static_cast<int>(refl.inv_d2[m] / s2_max * SHELLS));
            d[k] += rows.d[m];
            luzzati[k] += std::exp(-2 * M_PI * M_PI * sigma_x * sigma_x * refl.inv_d2[m]);
            n[k]++;
        }
        for (int k = 0; k < SHELLS; k++) {
            d[k] /= n[k];
            luzzati[k] /= n[k];
        }
    };
    XrayRowModel rows;
    rows.k_total.assign(refl.Size(), 1.0f);
    rows.fmask.assign(refl.Size(), {0, 0});
    for (float &s : refl.sigma)
        s = 0.01f;
    EstimateLikelihoodParameters(refl, f, rows);
    std::vector<double> d_clean, luzzati;
    shells(rows, d_clean, luzzati);
    for (int k = 0; k < SHELLS; k++) {
        CAPTURE(k, d_clean[k], luzzati[k]);
        CHECK(std::fabs(d_clean[k] - luzzati[k]) < 0.05);
    }
    CHECK(luzzati[SHELLS - 1] < 0.65);

    // a measurement error of half the rms amplitude of the outermost shell
    double fo2 = 0;
    int n = 0;
    for (size_t m = 0; m < refl.Size(); m++)
        if (refl.inv_d2[m] > s2_max * (SHELLS - 1) / SHELLS) {
            fo2 += double(refl.fo[m]) * refl.fo[m];
            n++;
        }
    for (float &s : refl.sigma)
        s = static_cast<float>(0.5 * std::sqrt(fo2 / n));
    EstimateLikelihoodParameters(refl, f, rows);
    std::vector<double> d_noisy;
    shells(rows, d_noisy, luzzati);
    CHECK(d_noisy[0] > d_clean[0] - 0.03);
    CHECK(d_noisy[SHELLS - 1] < d_clean[SHELLS - 1] - 0.1);
}

TEST_CASE("RefinementXray_gradient_CPU", "[Refinement]") {
    for (const char *group : {"P 1", "P 21 21 21", "C 1 2 1", "P 61 2 2"}) {
        CAPTURE(group);
        const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group);
        const gemmi::UnitCell cell = sg.crystal_system() == gemmi::CrystalSystem::Hexagonal
                                         ? gemmi::UnitCell(30, 30, 45, 90, 90, 120)
                                         : gemmi::UnitCell(28, 31, 34, 90, sg.number == 5 ? 104 : 90, 90);
        gemmi::Structure truth = random_crystal(cell, group, 240, 7);
        gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 8);
        RefinementReflections refl = PrepareRefinementReflections(asu_reflections(cell, sg, 2.2), cell, sg, 2.2);
        REQUIRE(refl.Size() > 500);
        auto cpu = MakeRefinementEngineCPU(cell, sg, 2.2, refl, 4);
        cpu->SetBlur(RefinementBlur(cell, sg, 2.2, 5.0));
        set_observations(*cpu, refl, truth.models[0], model);
        check_xray_gradient(*cpu, model, 0.03);
    }
}

#ifdef JFJOCH_USE_CUDA
TEST_CASE("RefinementXray_gradient_GPU", "[Refinement]") {
    if (get_gpu_count() == 0)
        return;
    for (const char *group : {"P 1", "P 21 21 21", "C 1 2 1", "P 61 2 2"}) {
        CAPTURE(group);
        const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group);
        const gemmi::UnitCell cell = sg.crystal_system() == gemmi::CrystalSystem::Hexagonal
                                         ? gemmi::UnitCell(30, 30, 45, 90, 90, 120)
                                         : gemmi::UnitCell(28, 31, 34, 90, sg.number == 5 ? 104 : 90, 90);
        gemmi::Structure truth = random_crystal(cell, group, 240, 7);
        gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 8);
        RefinementReflections refl = PrepareRefinementReflections(asu_reflections(cell, sg, 2.2), cell, sg, 2.2);
        auto cpu = MakeRefinementEngineCPU(cell, sg, 2.2, refl, 4);
        cpu->SetBlur(RefinementBlur(cell, sg, 2.2, 5.0));
        set_observations(*cpu, refl, truth.models[0], model);
        // the GPU engine sees the same observations: it is made after they are set
        auto gpu = MakeRefinementEngineGPU(0, cell, sg, 2.2, refl, model, 300.0, SIZE_MAX);
        REQUIRE(gpu);
        gpu->SetBlur(RefinementBlur(cell, sg, 2.2, 5.0));
        std::vector<std::complex<float>> f;
        cpu->Fcalc(model, f);
        XrayRowModel rows;
        rows.k_total.assign(refl.Size(), 1.0f);
        rows.fmask.assign(refl.Size(), {0, 0});
        EstimateLikelihoodParameters(refl, f, rows);
        cpu->SetRowModel(rows);
        gpu->SetRowModel(rows);

        std::vector<std::complex<float>> fg;
        gpu->Fcalc(model, fg);
        double num = 0, den = 0;
        for (size_t m = 0; m < f.size(); m++) {
            num += std::abs(f[m] - fg[m]);
            den += std::abs(f[m]);
        }
        CHECK(num / den < 1e-3);

        std::vector<std::array<double, 4>> gc, gg;
        const double tc = cpu->Evaluate(model, &gc);
        const double tg = gpu->Evaluate(model, &gg);
        CHECK(tg == Catch::Approx(tc).epsilon(1e-4));
        double max_g = 0;
        for (const auto &g : gc)
            for (int c = 0; c < 4; c++)
                max_g = std::max(max_g, std::fabs(g[c]));
        for (size_t i = 0; i < gc.size(); i++)
            for (int c = 0; c < 4; c++) {
                CAPTURE(i, c);
                CHECK(std::fabs(gg[i][c] - gc[i][c]) < 0.01 * max_g);
            }
        check_xray_gradient(*gpu, model, 0.05);
    }
}
#endif

TEST_CASE("RefinementMinimiser_quadratic", "[Refinement]") {
    // f = sum a_k (x_k - 1)^2 with curvatures spread over two decades
    const size_t n = 50;
    std::vector<double> a(n);
    for (size_t k = 0; k < n; k++)
        a[k] = std::pow(100.0, static_cast<double>(k) / (n - 1));
    const RefinementMinimiser::Target target = [&](const std::vector<double> &x, double &f, std::vector<double> &g) {
        f = 0;
        for (size_t k = 0; k < n; k++) {
            f += a[k] * (x[k] - 1) * (x[k] - 1);
            g[k] = 2 * a[k] * (x[k] - 1);
        }
        return true;
    };
    auto distance = [&](const std::vector<double> &x) {
        double d = 0;
        for (double v : x)
            d = std::max(d, std::fabs(v - 1));
        return d;
    };

    // the exact diagonal: one step
    RefinementMinimiser exact;
    std::vector<double> h(n);
    for (size_t k = 0; k < n; k++)
        h[k] = 2 * a[k];
    exact.SetDiagonal(h);
    std::vector<double> x(n, 0.0);
    exact.Minimise(target, x, 1, 0.0, 10.0);
    CHECK(distance(x) < 1e-9);

    // without it, the memory carried from one call to the next does what one long call does
    RefinementMinimiser split, whole;
    std::vector<double> xs(n, 0.0), xw(n, 0.0);
    for (int call = 0; call < 6; call++)
        split.Minimise(target, xs, 10, 0.0, 10.0);
    whole.Minimise(target, xw, 60, 0.0, 10.0);
    CHECK(distance(xw) < 1e-3);
    for (size_t k = 0; k < n; k++)
        CHECK(xs[k] == Catch::Approx(xw[k]).margin(1e-12));
}

TEST_CASE("Refinement_restrained_recovers_model", "[Refinement]") {
    const gemmi::UnitCell cell(36, 40, 44, 90, 90, 90);
    const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 21 21 21");
    gemmi::Structure truth = residue_crystal({"PHE", "TRP", "DA", "DG", "TYR", "HOH", "HOH", "HOH"}, cell,
                                             "P 21 21 21", 11);
    gemmi::Structure start = truth;
    start.models[0] = shaken(truth.models[0], 0.25, 5.0, 12);
    std::vector<MergedReflection> merged = asu_reflections(cell, sg, 1.8);
    {
        RefinementReflections refl = PrepareRefinementReflections(merged, cell, sg, 1.8 * 0.999);
        auto cpu = MakeRefinementEngineCPU(cell, sg, 1.8 * 0.999, refl, 4);
        cpu->SetBlur(RefinementBlur(cell, sg, 1.8, 5.0));
        std::vector<std::complex<float>> f;
        cpu->Fcalc(truth.models[0], f);
        std::map<gemmi::Miller, double> amplitude;
        for (size_t m = 0; m < refl.Size(); m++)
            amplitude[refl.hkl[m]] = std::abs(f[m]);
        std::mt19937 rng(3);
        std::normal_distribution<double> noise(0.0, 0.03);
        for (MergedReflection &r : merged) {
            const auto it = amplitude.find(gemmi::Miller{{r.h, r.k, r.l}});
            r.F = it == amplitude.end() ? NAN : static_cast<float>(it->second * (1 + noise(rng)));
            r.sigmaF = 0.03f * r.F + 0.01f;
        }
    }
    Logger logger("RefinementTest");
    for (bool gpu : {false, true}) {
        if (gpu && get_gpu_count() == 0)
            continue;
        CAPTURE(gpu);
        RefinementSettings settings;
        settings.gpu = gpu;
        settings.nthreads = 4;
        settings.weight_ladder = {0.0625};
        settings.adp_ladder = {4.0};
        settings.max_steps = 2;
        settings.final_ladder = {1.0};
        settings.riding_hydrogens = false;   // the data were made without them
        settings.solvent = true;
        const RefinementResult r = RefineModel(start, merged, cell, sg, 1.0, logger, settings);
        REQUIRE(r.ok);
        CHECK(r.on_gpu == gpu);
        CHECK(r.restrained_residues == 8);
        CHECK(r.r_free < 0.6 * r.r_free_start);
        CHECK(r.geometry.rmsd_bonds < r.geometry_start.rmsd_bonds);
        double rms = 0;
        std::vector<gemmi::Atom *> refined = atoms_of(r.model->models[0]);
        std::erase_if(refined, [](const gemmi::Atom *a) { return a->is_hydrogen(); });   // the riding hydrogens
        const std::vector<gemmi::Atom *> ref = atoms_of(truth.models[0]);
        REQUIRE(refined.size() == ref.size());   // the solvent build's placeholders not among them
        for (const gemmi::Chain &ch : r.model->models[0].chains)
            CHECK(ch.name != "~");
        for (size_t i = 0; i < ref.size(); i++)
            rms += refined[i]->pos.dist_sq(ref[i]->pos);
        rms = std::sqrt(rms / ref.size());
        CHECK(rms < 0.15);
        // the map coefficients as amplitude and phase: FWT and DELFWT non-negative, their phases in [0, 360)
        REQUIRE(r.maps);
        const auto &columns = r.maps->coefficients.columns;
        REQUIRE(columns[6].label == "FWT");
        REQUIRE(columns[8].label == "DELFWT");
        for (size_t c : {6, 8})
            for (size_t i = 0; i < columns[c].values.size(); i++) {
                CHECK(columns[c].values[i] >= 0.0f);
                CHECK(columns[c + 1].values[i] >= 0.0f);
                CHECK(columns[c + 1].values[i] < 360.0f);
            }

        // one trial at a time instead of three side by side (sharing a card's tables, mask and scale fit): the same
        // model, bit for bit
        RefinementSettings one = settings;
        one.max_trials_at_once = 1;
        const RefinementResult s = RefineModel(start, merged, cell, sg, 1.0, logger, one);
        REQUIRE(s.ok);
        CHECK(s.r_free == r.r_free);
        // the same trials kept, those run ahead where the slots take them left out
        CHECK(s.trials.size() == r.trials.size());
        CHECK(s.adp_trials.size() == r.adp_trials.size());
        CHECK(s.adp_model == r.adp_model);
        const std::vector<gemmi::Atom *> a_side = atoms_of(r.model->models[0]), a_one = atoms_of(s.model->models[0]);
        REQUIRE(a_one.size() == a_side.size());
        size_t same = 0;
        for (size_t i = 0; i < a_one.size(); i++)
            same += a_one[i]->pos.x == a_side[i]->pos.x && a_one[i]->pos.y == a_side[i]->pos.y &&
                    a_one[i]->pos.z == a_side[i]->pos.z && a_one[i]->b_iso == a_side[i]->b_iso;
        CHECK(same == a_one.size());
        // and the maps and what is read off them after it
        CHECK(s.mean_fom == r.mean_fom);
        CHECK(s.difference_map_rms == r.difference_map_rms);
        REQUIRE(s.blobs.size() == r.blobs.size());
        for (size_t i = 0; i < s.blobs.size(); i++) {
            CHECK(s.blobs[i].centroid.dist(r.blobs[i].centroid) == 0.0);
            CHECK(s.blobs[i].score == r.blobs[i].score);
        }
        REQUIRE(s.solvent);
        REQUIRE(r.solvent);
        CHECK(s.solvent->size() == r.solvent->size());

        // cut short before the first weight trial: the model as placed comes back, no trial chosen
        const std::atomic<bool> cancel{true};
        settings.cancel = &cancel;
        const RefinementResult c = RefineModel(start, merged, cell, sg, 1.0, logger, settings);
        REQUIRE(c.ok);
        CHECK(c.cancelled);
        CHECK_FALSE(c.restrained);
        CHECK(c.trials.empty());
        CHECK(c.model != nullptr);
        CHECK_FALSE(r.cancelled);
    }
}

// Two copies of a tRNA: with a slot for every ladder trial, the NCS rungs run ahead at both B weights and the
// intermediate TLS cuts with the others, those not needed left out; one trial at a time runs only what is needed.
// The same trials kept and the same model, bit for bit.
TEST_CASE("Refinement_trials_run_ahead", "[Refinement]") {
    gemmi::Structure truth = gemmi::read_structure_gz("../../tests/test_data/ligands/1ehz_psu.pdb.gz");
    truth.models.resize(1);
    gemmi::remove_hydrogens(truth.models[0]);
    truth.cell = gemmi::UnitCell(120, 120, 120, 90, 90, 90);
    truth.spacegroup_hm = "P 1";
    truth.setup_cell_images();
    gemmi::Chain copy = truth.models[0].chains.at(0);
    copy.name = "B";
    const double a = 0.7;
    const gemmi::Mat33 rot(std::cos(a), -std::sin(a), 0, std::sin(a), std::cos(a), 0, 0, 0, 1);
    for (gemmi::Residue &r : copy.residues)
        for (gemmi::Atom &at : r.atoms)
            at.pos = gemmi::Position(rot.multiply(at.pos)) + gemmi::Position(40, 10, -5);
    truth.models[0].chains.push_back(copy);
    truth.connections.clear();   // the LINK records name chain A only
    gemmi::Structure start = truth;
    start.models[0] = shaken(truth.models[0], 0.3, 5.0, 21);

    const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 1");
    const double d_min = 5.0;
    std::vector<MergedReflection> merged = asu_reflections(truth.cell, sg, d_min);
    {
        RefinementReflections refl = PrepareRefinementReflections(merged, truth.cell, sg, d_min * 0.999);
        auto cpu = MakeRefinementEngineCPU(truth.cell, sg, d_min * 0.999, refl, 8);
        cpu->SetBlur(RefinementBlur(truth.cell, sg, d_min, 5.0));
        std::vector<std::complex<float>> f;
        cpu->Fcalc(truth.models[0], f);
        std::map<gemmi::Miller, double> amplitude;
        for (size_t m = 0; m < refl.Size(); m++)
            amplitude[refl.hkl[m]] = std::abs(f[m]);
        for (MergedReflection &r : merged) {
            const auto it = amplitude.find(gemmi::Miller{{r.h, r.k, r.l}});
            r.F = it == amplitude.end() ? NAN : static_cast<float>(it->second);
            r.sigmaF = 0.03f * r.F + 0.01f;
        }
    }
    Logger logger("RefinementTest");
    RefinementSettings settings;
    settings.gpu = false;
    settings.nthreads = 8;
    settings.weight_ladder = {1.0 / 128, 1.0 / 16};
    settings.macro_cycles = 1;
    settings.max_steps = 1;
    settings.final_ladder = {1.0};
    settings.riding_hydrogens = false;
    settings.solvent = false;
    const RefinementResult ahead = RefineModel(start, merged, truth.cell, sg, 1.0, logger, settings);
    RefinementSettings one = settings;
    one.max_trials_at_once = 1;
    const RefinementResult needed = RefineModel(start, merged, truth.cell, sg, 1.0, logger, one);
    REQUIRE(ahead.ok);
    REQUIRE(needed.ok);
    CHECK(ahead.ncs_copies == 2);
    CHECK(ahead.trials.size() == needed.trials.size());
    CHECK(ahead.adp_trials.size() == needed.adp_trials.size());
    CHECK(ahead.adp_model == needed.adp_model);
    CHECK(ahead.r_free == needed.r_free);
    std::vector<gemmi::Atom *> x = atoms_of(ahead.model->models[0]), y = atoms_of(needed.model->models[0]);
    REQUIRE(x.size() == y.size());
    size_t same = 0;
    for (size_t i = 0; i < x.size(); i++)
        same += x[i]->pos.x == y[i]->pos.x && x[i]->pos.y == y[i]->pos.y && x[i]->pos.z == y[i]->pos.z &&
                x[i]->b_iso == y[i]->b_iso;
    CHECK(same == x.size());
}

// A K+ on the 4-fold axis of P4 at occupancy 0.25 (its four coinciding copies one ion), 0.2 A off the axis in the
// model as given, with a water 2.8 A from it: the topology puts the ion on the axis, its contact with the water is
// listed once (not once per copy) and none with itself; refinement keeps it on the axis at the occupancy given,
// while the rest of the model moves.
TEST_CASE("Refinement_special_position", "[Refinement]") {
    const gemmi::UnitCell cell(30, 30, 40, 90, 90, 90);
    const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 4");
    gemmi::Structure truth = residue_crystal({"PHE", "TRP", "TYR", "HOH", "HOH", "K"}, cell, "P 4", 5);
    std::vector<gemmi::Residue> &res = truth.models[0].chains[0].residues;
    const gemmi::Position axis(0, 0, 20);
    res[5].atoms[0].pos = axis;
    res[5].atoms[0].occ = 0.25f;
    res[4].atoms[0].pos = axis + gemmi::Position(2.8, 0, 0.3);
    gemmi::Structure start = truth;
    start.models[0] = shaken(truth.models[0], 0.25, 5.0, 6);
    gemmi::Atom &k_start = start.models[0].chains[0].residues[5].atoms[0];
    k_start.pos = axis + gemmi::Position(0.2, 0, 0.1);

    {
        gemmi::Structure st = start;
        const std::vector<SpecialPosition> sp = FindSpecialPositions(st);
        REQUIRE(sp.size() == 1);
        CHECK(sp[0].multiplicity == 4);
        const RefinementTopology topo = BuildRefinementTopology(st);
        const std::vector<gemmi::Atom *> atoms = atoms_of(st.models[0]);
        const int k = static_cast<int>(atoms.size()) - 1, water = k - 1;
        REQUIRE(topo.special.size() == 1);
        CHECK(topo.special[0].atom == k);
        CHECK(atoms[k]->pos.x == Catch::Approx(0).margin(1e-9));
        CHECK(atoms[k]->pos.y == Catch::Approx(0).margin(1e-9));
        CHECK(atoms[k]->pos.z == Catch::Approx(20.1));
        const RestraintSet rep = BuildRepulsions(topo, st.models[0], cell, sg, 1.5);
        int with_water = 0, with_itself = 0;
        for (const RestraintRecord &r : rep.records) {
            with_water += (r.atom[0] == k && r.atom[1] == water) || (r.atom[0] == water && r.atom[1] == k);
            with_itself += r.atom[0] == k && r.atom[1] == k;
        }
        CHECK(with_water == 1);
        CHECK(with_itself == 0);
        // the water's neighbours: the ion once, not four times
        gemmi::UnitCell images = cell;
        images.set_cell_images_from_spacegroup(&sg);
        gemmi::NeighborSearch ns(st.models[0], images, 5.0);
        ns.populate(false);
        int ion = 0;
        for (const ModelNeighbour &n : AtomsNear(ns, st.models[0], atoms[water]->pos, 4.0))
            ion += n.residue == 5;
        CHECK(ion == 1);
    }

    std::vector<MergedReflection> merged = asu_reflections(cell, sg, 1.8);
    {
        RefinementReflections refl = PrepareRefinementReflections(merged, cell, sg, 1.8 * 0.999);
        auto cpu = MakeRefinementEngineCPU(cell, sg, 1.8 * 0.999, refl, 4);
        cpu->SetBlur(RefinementBlur(cell, sg, 1.8, 5.0));
        std::vector<std::complex<float>> f;
        cpu->Fcalc(truth.models[0], f);
        std::map<gemmi::Miller, double> amplitude;
        for (size_t m = 0; m < refl.Size(); m++)
            amplitude[refl.hkl[m]] = std::abs(f[m]);
        for (MergedReflection &r : merged) {
            const auto it = amplitude.find(gemmi::Miller{{r.h, r.k, r.l}});
            r.F = it == amplitude.end() ? NAN : static_cast<float>(it->second);
            r.sigmaF = 0.03f * r.F + 0.01f;
        }
    }
    Logger logger("RefinementTest");
    RefinementSettings settings;
    settings.gpu = false;
    settings.nthreads = 4;
    settings.weight_ladder = {0.0625};
    settings.adp_ladder = {4.0};
    settings.max_steps = 2;
    settings.final_ladder = {1.0};
    settings.riding_hydrogens = false;
    settings.solvent = false;
    const RefinementResult r = RefineModel(start, merged, cell, sg, 1.0, logger, settings);
    REQUIRE(r.ok);
    REQUIRE(r.restrained);
    const gemmi::Atom &k = r.model->models[0].chains[0].residues[5].atoms[0];
    CHECK(std::hypot(k.pos.x, k.pos.y) < 1e-6);
    CHECK(k.pos.z == Catch::Approx(20).margin(0.15));
    CHECK(k.occ == 0.25f);
    CHECK(r.r_free < 0.6 * r.r_free_start);
}

// Single-atom ions scatter as the ions they are (10 electrons for Mg2+, 18 for Cl-), every other atom as the
// neutral atom whatever charge the file gives it; the CPU and the GPU engines score such a model alike.
TEST_CASE("Refinement_ion_form_factors", "[Refinement]") {
    CHECK(IonFormFactor<float>(gemmi::El::Mg, 2).calculate_sf(0.0f) == Catch::Approx(10.0).margin(0.02));
    CHECK(IonFormFactor<float>(gemmi::El::Mg, 0).calculate_sf(0.0f) == Catch::Approx(12.0).margin(0.02));
    CHECK(IonFormFactor<double>(gemmi::El::Cl, -1).calculate_sf(0.0) == Catch::Approx(18.0).margin(0.02));
    CHECK(IonFormFactor<float>(gemmi::El::N, 1).calculate_sf(0.0f) == Catch::Approx(7.0).margin(0.02));   // none tabulated

    const gemmi::UnitCell cell(30, 32, 34, 90, 90, 90);
    const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 21 21 21");
    gemmi::Structure st = residue_crystal({"LYS", "TRP", "MG", "CL", "HOH"}, cell, "P 21 21 21", 3);
    for (gemmi::Atom &a : st.models[0].chains[0].residues[0].atoms)
        if (a.name == "NZ")
            a.charge = 1;   // as a file may give it
    std::vector<MergedReflection> merged = asu_reflections(cell, sg, 2.0);
    std::mt19937 rng(4);
    std::uniform_real_distribution<float> f(10.0f, 100.0f);
    for (MergedReflection &r : merged) {
        r.F = f(rng);
        r.sigmaF = 1.0f;
    }
    Logger logger("RefinementTest");
    RefinementSettings settings;
    settings.score_only = true;
    settings.nthreads = 4;
    settings.gpu = false;
    const RefinementResult cpu = RefineModel(st, merged, cell, sg, 1.0, logger, settings);
    REQUIRE(cpu.ok);
    for (const gemmi::Residue &r : cpu.model->models[0].chains[0].residues)
        for (const gemmi::Atom &a : r.atoms)
            CHECK(static_cast<int>(a.charge) == (r.name == "MG" ? 2 : r.name == "CL" ? -1 : 0));
    if (get_gpu_count() > 0) {
        settings.gpu = true;
        const RefinementResult gpu = RefineModel(st, merged, cell, sg, 1.0, logger, settings);
        REQUIRE(gpu.ok);
        CHECK(gpu.on_gpu);
        CHECK(gpu.r_work_start == Catch::Approx(cpu.r_work_start).margin(1e-4));
    }
}

// Two copies of one RNA chain in the asymmetric unit, the second turned and moved, a fragment of it and the chain in
// another conformation: the two found as copies of one molecule, the fragment and the other conformation not, held
// to each other by local distance restraints whose gradient is right on the CPU and the GPU, and let go where the
// copies really differ. A single copy has none.
TEST_CASE("RefinementNcs_two_copies", "[Refinement]") {
    gemmi::Structure st = gemmi::read_structure_gz("../../tests/test_data/ligands/1ehz_psu.pdb.gz");
    st.models.resize(1);
    gemmi::remove_hydrogens(st.models[0]);
    st.cell = gemmi::UnitCell(120, 120, 120, 90, 90, 90);
    st.spacegroup_hm = "P 1";
    st.setup_cell_images();
    {
        gemmi::Structure single = st;
        const RefinementTopology topo = BuildRefinementTopology(single);
        CHECK(FindNcsCopies(single.models[0], topo).Empty());
    }
    gemmi::Chain copy = st.models[0].chains.at(0);
    copy.name = "B";
    const double a = 0.7;
    const gemmi::Mat33 rot(std::cos(a), -std::sin(a), 0, std::sin(a), std::cos(a), 0, 0, 0, 1);
    for (gemmi::Residue &r : copy.residues)
        for (gemmi::Atom &at : r.atoms)
            at.pos = gemmi::Position(rot.multiply(at.pos)) + gemmi::Position(40, 10, -5);
    st.models[0].chains.push_back(copy);
    // and a fragment of the molecule, its first 18 residues: not a copy of the whole
    gemmi::Chain fragment = st.models[0].chains.at(0);
    fragment.name = "C";
    fragment.residues.resize(18);
    for (gemmi::Residue &r : fragment.residues)
        for (gemmi::Atom &at : r.atoms)
            at.pos += gemmi::Position(-30, 40, 20);
    st.models[0].chains.push_back(fragment);
    // and the whole molecule in another conformation, each residue moved 3 A its own way: the same sequence, not a copy
    gemmi::Chain other = st.models[0].chains.at(0);
    other.name = "D";
    std::mt19937 shift_rng(11);
    std::uniform_real_distribution<double> shift(-1.0, 1.0);
    for (gemmi::Residue &r : other.residues) {
        const gemmi::Vec3 v = gemmi::Vec3(shift(shift_rng), shift(shift_rng), shift(shift_rng)).normalized() * 3.0;
        for (gemmi::Atom &at : r.atoms)
            at.pos += gemmi::Position(v) + gemmi::Position(30, -40, 25);
    }
    st.models[0].chains.push_back(other);
    st.connections.clear();   // the LINK records name chain A only
    const RefinementTopology topo = BuildRefinementTopology(st);
    const NcsCopies ncs = FindNcsCopies(st.models[0], topo);
    REQUIRE_FALSE(ncs.Empty());
    CHECK(ncs.groups == 1);
    CHECK(ncs.copies == 2);
    CHECK(ncs.description == "A=B");
    CHECK(ncs.anchors.size() == 1);
    CHECK(ncs.anchors[0].size() > 25);
    CAPTURE(ncs.distances.records.size());
    CHECK(ncs.distances.records.size() > 1000);

    std::vector<gemmi::Atom *> atoms = atoms_of(st.models[0]);
    std::vector<std::array<double, 4>> xyzb;
    for (gemmi::Atom *at : atoms)
        xyzb.push_back({at->pos.x, at->pos.y, at->pos.z, at->b_iso});
    RestraintSet set = NcsRestraintsAt(ncs, 0.1);
    set.Finalize(xyzb.size());
    // the same molecule twice: nothing to restrain
    CHECK(EvaluateRestraints(set, xyzb, nullptr) == Catch::Approx(0.0).margin(1e-6));
    CHECK(NcsRmsd(ncs, xyzb) < 1e-3);
    CHECK(NcsReleasedResidues(ncs, xyzb) == 0);

    // the second copy shaken: the gradient against central differences
    std::mt19937 rng(7);
    std::normal_distribution<double> d(0, 0.1);
    const size_t half = gemmi::count_atom_sites(st.models[0].chains[0]);
    for (size_t i = half; i < 2 * half; i++)
        for (int c = 0; c < 3; c++)
            xyzb[i][c] += d(rng);
    std::vector<std::array<double, 4>> grad(xyzb.size(), {0, 0, 0, 0});
    const double e = EvaluateRestraints(set, xyzb, &grad);
    CHECK(e > 0);
    const double h = 1e-6;
    for (size_t i = 0; i < xyzb.size(); i += 17)
        for (int c = 0; c < 3; c++) {
            auto x = xyzb;
            x[i][c] += h;
            const double fp = EvaluateRestraints(set, x, nullptr);
            x[i][c] -= 2 * h;
            const double fm = EvaluateRestraints(set, x, nullptr);
            CAPTURE(i, c);
            CHECK(grad[i][c] == Catch::Approx((fp - fm) / (2 * h)).epsilon(1e-4).margin(1e-3));
        }

#ifdef JFJOCH_USE_CUDA
    if (get_gpu_count() > 0) {
        const gemmi::SpaceGroup &p1 = *gemmi::find_spacegroup_by_name("P 1");
        const gemmi::UnitCell wide_cell(40, 40, 40, 90, 90, 90);
        std::vector<MergedReflection> merged = asu_reflections(wide_cell, p1, 3.0);
        RefinementReflections refl = PrepareRefinementReflections(merged, wide_cell, p1, 3.0);
        auto gpu = MakeRefinementEngineGPU(0, wide_cell, p1, 3.0, refl, st.models[0], 100.0, SIZE_MAX);
        REQUIRE(gpu);
        gpu->SetRestraints(0, set);
        std::vector<std::array<double, 4>> ggpu;
        const double egpu = gpu->EvaluateRestraints(0, xyzb, &ggpu);
        CHECK(egpu == Catch::Approx(e).epsilon(1e-3));
        for (size_t i = 0; i < xyzb.size(); i++)
            for (int c = 0; c < 3; c++)
                CHECK(ggpu[i][c] == Catch::Approx(grad[i][c]).epsilon(2e-3).margin(2e-2));
    }
#endif

    // one residue of the second copy (its 16th) distorted, each atom moved by 1.5 A in a direction of its own: its
    // restraints let go, the others hold
    size_t first = gemmi::count_atom_sites(st.models[0].chains[0]);
    for (int k = 0; k < 15; k++)
        first += st.models[0].chains[1].residues[k].atoms.size();
    std::vector<std::array<double, 4>> moved_xyzb = xyzb;
    std::uniform_real_distribution<double> u(-1.0, 1.0);
    for (size_t i = first; i < first + st.models[0].chains[1].residues[15].atoms.size(); i++) {
        const gemmi::Vec3 v = gemmi::Vec3(u(rng), u(rng), u(rng)).normalized() * 1.5;
        moved_xyzb[i][0] += v.x;
        moved_xyzb[i][1] += v.y;
        moved_xyzb[i][2] += v.z;
    }
    CHECK(NcsReleasedResidues(ncs, xyzb) == 0);
    CHECK(NcsReleasedResidues(ncs, moved_xyzb) >= 1);
    CHECK(NcsReleasedResidues(ncs, moved_xyzb) <= 3);
}

namespace {

// Every other atom anisotropic, its dev random but traceless, up to a third of its B.
std::vector<AtomAniso> random_aniso(const gemmi::Model &model, unsigned seed) {
    std::mt19937 rng(seed);
    std::uniform_real_distribution<double> d(-1.0, 1.0);
    std::vector<AtomAniso> out;
    int i = 0;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            for (const gemmi::Atom &a : r.atoms) {
                AtomAniso an;
                an.on = (i++ % 2) == 0;
                if (an.on) {
                    const double u = a.b_iso / gemmi::u_to_b() / 3.0;
                    double v[6];
                    for (double &x : v)
                        x = u * d(rng);
                    const double t = (v[0] + v[1] + v[2]) / 3.0;
                    an.dev = {static_cast<float>(v[0] - t), static_cast<float>(v[1] - t), static_cast<float>(v[2] - t),
                              static_cast<float>(v[3] / 2), static_cast<float>(v[4] / 2), static_cast<float>(v[5] / 2)};
                }
                out.push_back(an);
            }
    return out;
}

// The engine's U gradient of a few anisotropic atoms against central differences of its own target.
void check_u_gradient(RefinementEngine &engine, const gemmi::Model &model, std::vector<AtomAniso> &aniso,
                      double tolerance) {
    std::vector<std::array<double, 4>> grad;
    engine.SetAniso(&aniso);
    engine.Evaluate(model, &grad);
    const std::vector<std::array<double, 6>> gu = engine.UGradient();
    REQUIRE(gu.size() == aniso.size());
    double max_g = 0;
    for (size_t i = 0; i < gu.size(); i++)
        for (int c = 0; c < 6; c++)
            max_g = std::max(max_g, std::fabs(gu[i][c]));
    const double h = 2e-3;   // A^2
    for (size_t i = 0; i < aniso.size(); i += 2 * (aniso.size() / 14)) {
        REQUIRE(aniso[i].on);
        for (int c = 0; c < 6; c++) {
            const float v0 = aniso[i].dev[c];
            aniso[i].dev[c] = static_cast<float>(v0 + h);
            const double fp = engine.Evaluate(model, nullptr);
            aniso[i].dev[c] = static_cast<float>(v0 - h);
            const double fm = engine.Evaluate(model, nullptr);
            aniso[i].dev[c] = v0;
            const double fd = (fp - fm) / (2 * h);
            CAPTURE(i, c, fd, gu[i][c], max_g);
            CHECK(std::fabs(fd - gu[i][c]) < tolerance * std::max(std::fabs(fd), 0.05 * max_g));
        }
    }
}

}  // namespace

TEST_CASE("RefinementXray_aniso_gradient_CPU", "[Refinement]") {
    for (const char *group : {"P 1", "P 21 21 21"}) {
        CAPTURE(group);
        const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group);
        const gemmi::UnitCell cell(28, 31, 34, 90, 90, 90);
        gemmi::Structure truth = random_crystal(cell, group, 240, 7);
        gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 8);
        RefinementReflections refl = PrepareRefinementReflections(asu_reflections(cell, sg, 2.0), cell, sg, 2.0);
        auto cpu = MakeRefinementEngineCPU(cell, sg, 2.0, refl, 4);
        cpu->SetBlur(RefinementBlur(cell, sg, 2.0, 5.0));
        set_observations(*cpu, refl, truth.models[0], model);
        std::vector<AtomAniso> aniso = random_aniso(model, 9);
        check_u_gradient(*cpu, model, aniso, 0.06);
        // the positions and the B along the isotropic direction, the shapes held
        check_xray_gradient(*cpu, model, 0.03);
    }
}

#ifdef JFJOCH_USE_CUDA
TEST_CASE("RefinementXray_aniso_gradient_GPU", "[Refinement]") {
    if (get_gpu_count() == 0)
        return;
    const char *group = "P 21 21 21";
    const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group);
    const gemmi::UnitCell cell(28, 31, 34, 90, 90, 90);
    gemmi::Structure truth = random_crystal(cell, group, 240, 7);
    gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 8);
    RefinementReflections refl = PrepareRefinementReflections(asu_reflections(cell, sg, 2.0), cell, sg, 2.0);
    auto cpu = MakeRefinementEngineCPU(cell, sg, 2.0, refl, 4);
    cpu->SetBlur(RefinementBlur(cell, sg, 2.0, 5.0));
    set_observations(*cpu, refl, truth.models[0], model);
    auto gpu = MakeRefinementEngineGPU(0, cell, sg, 2.0, refl, model, 300.0, SIZE_MAX);
    REQUIRE(gpu);
    gpu->SetBlur(RefinementBlur(cell, sg, 2.0, 5.0));
    std::vector<std::complex<float>> f;
    std::vector<AtomAniso> aniso = random_aniso(model, 9);
    cpu->SetAniso(&aniso);
    gpu->SetAniso(&aniso);
    cpu->Fcalc(model, f);
    XrayRowModel rows;
    rows.k_total.assign(refl.Size(), 1.0f);
    rows.fmask.assign(refl.Size(), {0, 0});
    EstimateLikelihoodParameters(refl, f, rows);
    cpu->SetRowModel(rows);
    gpu->SetRowModel(rows);
    std::vector<std::complex<float>> fg;
    gpu->Fcalc(model, fg);
    double num = 0, den = 0;
    for (size_t m = 0; m < f.size(); m++) {
        num += std::abs(f[m] - fg[m]);
        den += std::abs(f[m]);
    }
    CHECK(num / den < 1e-3);
    std::vector<std::array<double, 4>> gc, gg;
    cpu->Evaluate(model, &gc);
    const std::vector<std::array<double, 6>> uc = cpu->UGradient();
    gpu->Evaluate(model, &gg);
    const std::vector<std::array<double, 6>> ug = gpu->UGradient();
    REQUIRE(ug.size() == uc.size());
    double max_g = 0, max_u = 0;
    for (size_t i = 0; i < gc.size(); i++) {
        for (int c = 0; c < 4; c++)
            max_g = std::max(max_g, std::fabs(gc[i][c]));
        for (int c = 0; c < 6; c++)
            max_u = std::max(max_u, std::fabs(uc[i][c]));
    }
    for (size_t i = 0; i < gc.size(); i++) {
        CAPTURE(i);
        for (int c = 0; c < 4; c++)
            CHECK(std::fabs(gg[i][c] - gc[i][c]) < 0.01 * max_g);
        for (int c = 0; c < 6; c++)
            CHECK(std::fabs(ug[i][c] - uc[i][c]) < 0.01 * max_u);
    }
    check_u_gradient(*gpu, model, aniso, 0.08);
}

// The GPU engine's evaluation with the gradient, timed on a protein-sized case: about 5500 atoms in the asymmetric
// unit of a P 21 21 21 cell at 1.6 A, the density of a protein's atoms with their hydrogens. Run by hand, under a
// profiler for the time of each kernel.
TEST_CASE("RefinementGPU_Benchmark", "[.][refine_bench]") {
    if (get_gpu_count() == 0)
        return;
    const char *group = "P 21 21 21";
    const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group);
    const gemmi::UnitCell cell(60, 64, 69, 90, 90, 90);
    const double d_min = 1.6;
    gemmi::Structure truth = random_crystal(cell, group, 5500, 7);
    gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 8);
    RefinementReflections refl = PrepareRefinementReflections(asu_reflections(cell, sg, d_min), cell, sg, d_min);
    auto gpu = MakeRefinementEngineGPU(0, cell, sg, d_min, refl, model, 300.0, SIZE_MAX);
    REQUIRE(gpu);
    gpu->SetBlur(RefinementBlur(cell, sg, d_min, 5.0));
    set_observations(*gpu, refl, truth.models[0], model);
    std::vector<std::array<double, 4>> g;
    gpu->Evaluate(model, &g);
    const int n = 200;
    const auto start = std::chrono::steady_clock::now();
    double t = 0;
    for (int i = 0; i < n; i++)
        t += gpu->Evaluate(model, &g);
    const double ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - start).count();
    printf("RefinementGPU_Benchmark: %zu atoms, %zu reflections, %.3f ms per evaluation (target %.6g)\n",
           atoms_of(model).size(), refl.Size(), ms / n, t / n);
}
#endif

TEST_CASE("RefinementADP_tls_gradient", "[Refinement]") {
    std::mt19937 rng(4);
    std::uniform_real_distribution<double> d(-1.0, 1.0);
    gemmi::Model model(1);
    gemmi::Chain chain("A");
    gemmi::Residue res;
    for (int i = 0; i < 30; i++) {
        gemmi::Atom a;
        a.pos = gemmi::Position(10 + 8 * d(rng), -3 + 8 * d(rng), 5 + 8 * d(rng));
        res.atoms.push_back(a);
    }
    chain.residues.push_back(res);
    model.chains.push_back(chain);
    std::vector<gemmi::Atom *> atoms = atoms_of(model);
    TlsBody body;
    body.origin = gemmi::Position(9, -2, 4);
    std::array<double, TlsBody::PARAMETERS> p{};
    for (double &x : p)
        x = 0.01 * d(rng);
    body.Set(p);
    for (int i = 0; i < 30; i++)
        body.atoms.push_back(i);
    // U is linear in the parameters: U(p) is sum p_k dU/dp_k
    for (gemmi::Atom *a : atoms) {
        const gemmi::SMat33<double> u = body.U(a->pos);
        const auto dk = body.Derivatives(a->pos);
        const double uu[6] = {u.u11, u.u22, u.u33, u.u12, u.u13, u.u23};
        for (int c = 0; c < 6; c++) {
            double sum = 0;
            for (int k = 0; k < TlsBody::PARAMETERS; k++)
                sum += p[k] * dk[k][c];
            CHECK(sum == Catch::Approx(uu[c]).margin(1e-12));
        }
    }
    // the chain rule in matrix form against the sum over dU/dp
    std::vector<std::array<double, 6>> gu(atoms.size());
    for (auto &g : gu)
        for (double &x : g)
            x = d(rng);
    const std::array<double, TlsBody::PARAMETERS> g = body.Gradient(atoms, gu);
    for (int k = 0; k < TlsBody::PARAMETERS; k++) {
        double sum = 0;
        for (size_t i = 0; i < atoms.size(); i++) {
            const auto dk = body.Derivatives(atoms[i]->pos);
            for (int c = 0; c < 6; c++)
                sum += gu[i][c] * dk[k][c];
        }
        CAPTURE(k);
        CHECK(g[k] == Catch::Approx(sum).margin(1e-9));
    }
    // and through the metadata in degrees and back
    gemmi::Model m2 = model;
    const std::vector<gemmi::TlsGroup> groups = TlsBodiesToGroups(m2, {body});
    REQUIRE(groups.size() == 1);
    const std::vector<TlsBody> back = TlsBodiesFromModel(m2, groups);
    REQUIRE(back.size() == 1);
    CHECK(back[0].atoms.size() == 30);
    const gemmi::SMat33<double> u0 = body.U(atoms[3]->pos), u1 = back[0].U(atoms[3]->pos);
    CHECK(u1.u12 == Catch::Approx(u0.u12).margin(1e-12));
    CHECK(u1.u33 == Catch::Approx(u0.u33).margin(1e-12));
}

TEST_CASE("RefinementADP_u_restraints_gradient", "[Refinement]") {
    std::mt19937 rng(5);
    std::uniform_real_distribution<double> d(-1.0, 1.0);
    const int n = 6;
    std::vector<gemmi::Position> pos;
    std::vector<std::array<double, 6>> u(n);
    for (int i = 0; i < n; i++) {
        pos.emplace_back(1.5 * i + 0.2 * d(rng), 0.3 * d(rng), 0.3 * d(rng));
        u[i] = {0.2 + 0.05 * d(rng), 0.2 + 0.05 * d(rng), 0.2 + 0.05 * d(rng), 0.03 * d(rng), 0.03 * d(rng), 0.03 * d(rng)};
    }
    const URestraints r = MakeURestraints({{0, 1}, {1, 2}, {0, 2}, {3, 4}}, {1, 1, 1, 1, 1, 1});
    CHECK(r.pairs.size() == 4);
    REQUIRE(r.isotropic == std::vector<int>{5});
    std::vector<std::array<double, 6>> g(n, {0, 0, 0, 0, 0, 0});
    EvaluateURestraints(r, pos, u, &g);
    const double h = 1e-6;
    for (int i = 0; i < n; i++)
        for (int c = 0; c < 6; c++) {
            std::vector<std::array<double, 6>> up = u, um = u;
            up[i][c] += h;
            um[i][c] -= h;
            const double fd = (EvaluateURestraints(r, pos, up, nullptr) - EvaluateURestraints(r, pos, um, nullptr)) / (2 * h);
            CAPTURE(i, c);
            CHECK(g[i][c] == Catch::Approx(fd).epsilon(1e-5).margin(1e-3));
        }
    // a rigid pair: the same U along the bond, however different across it, costs no rigid-bond term
    const URestraints one = MakeURestraints({{0, 1}}, {1, 1});
    const std::vector<gemmi::Position> line = {{0, 0, 0}, {0, 0, 1.5}};
    const std::vector<std::array<double, 6>> across = {{0.1, 0.3, 0.2, 0, 0, 0}, {0.3, 0.1, 0.2, 0.05, 0, 0}};
    const double sim = (0.2 * 0.2 + 0.2 * 0.2 + 2 * 0.05 * 0.05) / (URestraints::SIGMA_SIMILAR * URestraints::SIGMA_SIMILAR);
    CHECK(EvaluateURestraints(one, line, across, nullptr) == Catch::Approx(sim));
}

TEST_CASE("RefinementADP_tls_segments", "[Refinement]") {
    // a chain of 60 residues along x whose B follows one quadratic on the first 25 residues and another after
    std::vector<std::vector<int>> residues;
    std::vector<gemmi::Position> pos;
    std::vector<double> b;
    for (int r = 0; r < 60; r++) {
        std::vector<int> ra;
        for (int a = 0; a < 5; a++) {
            const gemmi::Position p(3.8 * r, 2.0 * std::sin(a + r), 2.0 * std::cos(a * r));
            ra.push_back(static_cast<int>(pos.size()));
            pos.push_back(p);
            const double y = p.y, z = p.z;
            b.push_back(r < 25 ? 20 + 0.5 * (y * y + z * z) : 60 + 3.0 * (y * y) - 0.02 * p.x);
        }
        residues.push_back(ra);
    }
    CHECK(TlsSegments(residues, pos, b, 1, 10) == std::vector<int>{0});
    CHECK(TlsSegments(residues, pos, b, 2, 10) == std::vector<int>{0, 25});
    // too short for three segments of 25
    CHECK(TlsSegments(residues, pos, b, 3, 25).size() == 2);
}

namespace {

// A twinned test case in P 4, whose lattice adds the twin law: the reflections of the group's ASU to 2.2 A with
// every fourth of the zone l = 1 left out, so that some twin mates are rows of their own, and the free set drawn
// on the lattice's holohedry (AssignRfreeFlags) so that no twin pair is split.
struct TwinCase {
    gemmi::UnitCell cell{30, 30, 40, 90, 90, 90};
    const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 4");
    gemmi::Op law;
    std::vector<MergedReflection> merged;
    RefinementReflections refl;

    TwinCase() {
        const std::vector<gemmi::Op> laws = RefinementTwinLaws(cell, sg);
        REQUIRE(laws.size() == 1);
        law = laws[0];
        merged = asu_reflections(cell, sg, 2.2);
        int i = 0;
        std::erase_if(merged, [&](const MergedReflection &r) { return r.l == 1 && i++ % 4 == 0; });
        AssignRfreeFlags(merged, sg, 0.05, 0,
                         UnitCell{static_cast<float>(cell.a), static_cast<float>(cell.b), static_cast<float>(cell.c),
                                  90, 90, 90});
        refl = PrepareRefinementReflections(merged, cell, sg, 2.2);
        CHECK(SetTwinLaw(refl, cell, sg, 2.2, law) == 0);
    }
};

}  // namespace

TEST_CASE("RefinementTwin_laws_and_mates", "[Refinement]") {
    TwinCase t;
    CHECK(RefinementTwinLaws(t.cell, *gemmi::find_spacegroup_by_name("P 4 2 2")).empty());
    REQUIRE(t.refl.Twinned());
    CHECK(!t.refl.mate_hkl.empty());
    const gemmi::GroupOps gops = t.sg.operations();
    const gemmi::ReciprocalAsu asu(&t.sg);
    const std::vector<gemmi::Miller> rows = t.refl.RowHkl();
    for (size_t m = 0; m < t.refl.Rows(); m++) {
        const int j = t.refl.mate[m];
        CAPTURE(m, j);
        REQUIRE(j >= 0);
        CHECK(t.refl.mate[j] == static_cast<int>(m));   // an involution
        if (m < t.refl.Size())
            CHECK(rows[j] == asu.to_asu(t.law.apply_to_hkl(rows[m]), gops).first);
    }
    // a free set drawn on the merging group alone splits twin pairs
    RefinementReflections split = t.refl;
    for (size_t m = 0; m < split.Size(); m++)
        split.free[m] = m % 20 == 0;
    CHECK(SetTwinLaw(split, t.cell, t.sg, 2.2, t.law) > 0);
}

TEST_CASE("RefinementTwin_rule_with_data_evidence", "[Refinement]") {
    // the null alone: 3 standard errors
    CHECK(AdoptTwinLaw(0.1, 3.0, false));
    CHECK_FALSE(AdoptTwinLaw(0.1, 2.9, false));
    // the data read twinned as well: 2 standard errors
    CHECK(AdoptTwinLaw(0.1, 2.0, true));
    CHECK(AdoptTwinLaw(0.1, 2.5, true));
    CHECK_FALSE(AdoptTwinLaw(0.1, 1.9, true));
    // no fraction, no twin, whatever the evidence
    CHECK_FALSE(AdoptTwinLaw(0.0, 10.0, true));
}

TEST_CASE("RefinementTwin_gradient", "[Refinement]") {
    TwinCase t;
    gemmi::Structure truth = random_crystal(t.cell, "P 4", 240, 17);
    const gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 18);
    auto cpu = MakeRefinementEngineCPU(t.cell, t.sg, 2.2, t.refl, 4);
    cpu->SetBlur(RefinementBlur(t.cell, t.sg, 2.2, 5.0));
    XrayRowModel rows;
    rows.k_total.assign(t.refl.Size(), 1.0f);
    rows.fmask.assign(t.refl.Rows(), {0, 0});
    rows.twin_fraction = 0.3f;
    std::vector<std::complex<float>> f;
    cpu->Fcalc(truth.models[0], f);
    REQUIRE(f.size() == t.refl.Rows());
    const std::vector<std::complex<float>> twinned = TwinnedModel(t.refl, rows, f);
    for (size_t m = 0; m < t.refl.Size(); m++)
        t.refl.fo[m] = std::abs(twinned[m]);
    // the fraction comes back from the data
    CHECK(FitTwinFraction(t.refl, rows, f) == Catch::Approx(0.3).margin(0.01));
    cpu->Fcalc(model, f);
    EstimateLikelihoodParameters(t.refl, TwinnedModel(t.refl, rows, f), rows);
    cpu->SetRowModel(rows);
    // the CPU engine kept a reference to t.refl and reads the observations set above
    check_xray_gradient(*cpu, model, 0.03);
#ifdef JFJOCH_USE_CUDA
    if (get_gpu_count() > 0) {
        auto gpu = MakeRefinementEngineGPU(0, t.cell, t.sg, 2.2, t.refl, model, 300.0, SIZE_MAX);
        REQUIRE(gpu);
        gpu->SetBlur(RefinementBlur(t.cell, t.sg, 2.2, 5.0));
        gpu->SetRowModel(rows);
        std::vector<std::array<double, 4>> gc, gg;
        const double tc = cpu->Evaluate(model, &gc);
        const double tg = gpu->Evaluate(model, &gg);
        CHECK(tg == Catch::Approx(tc).epsilon(1e-4));
        double max_g = 0;
        for (const auto &g : gc)
            for (int c = 0; c < 4; c++)
                max_g = std::max(max_g, std::fabs(g[c]));
        for (size_t i = 0; i < gc.size(); i++)
            for (int c = 0; c < 4; c++) {
                CAPTURE(i, c);
                CHECK(std::fabs(gg[i][c] - gc[i][c]) < 0.01 * max_g);
            }
    }
#endif
}

TEST_CASE("Refinement_twin_decided", "[Refinement]") {
    TwinCase t;
    gemmi::Structure truth = residue_crystal({"PHE", "TRP", "DA", "DG", "TYR", "HOH", "HOH", "HOH"}, t.cell, "P 4", 21);
    gemmi::Structure start = truth;
    start.models[0] = shaken(truth.models[0], 0.2, 4.0, 22);
    std::vector<std::complex<float>> f;
    {
        RefinementReflections refl = t.refl;
        auto cpu = MakeRefinementEngineCPU(t.cell, t.sg, 2.2, refl, 4);
        cpu->SetBlur(RefinementBlur(t.cell, t.sg, 2.2, 5.0));
        cpu->Fcalc(truth.models[0], f);
    }
    Logger logger("RefinementTest");
    RefinementSettings settings;
    settings.gpu = false;
    settings.nthreads = 4;
    settings.weight_ladder = {0.0625};
    settings.adp_ladder = {0.0};
    settings.max_steps = 1;
    settings.final_ladder = {1.0};
    settings.riding_hydrogens = false;
    settings.solvent = false;
    for (double fraction : {0.0, 0.3}) {
        CAPTURE(fraction);
        // the observations as twinned at `fraction`, with 3% noise
        XrayRowModel rows;
        rows.k_total.assign(t.refl.Size(), 1.0f);
        rows.fmask.assign(t.refl.Rows(), {0, 0});
        rows.twin_fraction = static_cast<float>(fraction);
        const std::vector<std::complex<float>> twinned = TwinnedModel(t.refl, rows, f);
        std::map<gemmi::Miller, double> amplitude;
        for (size_t m = 0; m < t.refl.Size(); m++)
            amplitude[t.refl.hkl[m]] = std::abs(twinned[m]);
        std::vector<MergedReflection> merged = t.merged;
        std::mt19937 rng(5);
        std::normal_distribution<double> noise(0.0, 0.03);
        for (MergedReflection &r : merged) {
            const auto it = amplitude.find(gemmi::Miller{{r.h, r.k, r.l}});
            r.F = it == amplitude.end() ? NAN : static_cast<float>(it->second * (1 + noise(rng)));
            r.sigmaF = 0.03f * r.F + 0.01f;
        }
        const RefinementResult r = RefineModel(start, merged, t.cell, t.sg, 1.0, logger, settings);
        REQUIRE(r.ok);
        CHECK(r.twin_law == t.law.as_hkl().triplet('h'));
        if (fraction > 0) {
            CHECK(r.twin_decision == "TWINNED");
            CHECK(r.twinned);
            CHECK(r.twin_fraction == Catch::Approx(fraction).margin(0.03));
            CHECK(r.r_free < r.untwinned_r_free - 0.02);
            CHECK(r.twin_r_free < r.twin_r_free_null - 0.02);
            CHECK(r.twin_null_z >= 3.0);
        } else {
            CHECK((r.twin_decision == "ZERO_FRACTION" || r.twin_null_z < 3.0));
            CHECK(r.twin_decision != "TWINNED");
            CHECK_FALSE(r.twinned);
            // the untwinned refinement, as it is without the twin hypothesis
            RefinementSettings plain = settings;
            plain.twin = false;
            const RefinementResult u = RefineModel(start, merged, t.cell, t.sg, 1.0, logger, plain);
            CHECK(u.r_free == r.r_free);
            CHECK(u.r_work == r.r_work);
        }
        // scored only, the model as given is tested against the twin law by the same rule
        RefinementSettings score = settings;
        score.score_only = true;
        const RefinementResult s = RefineModel(truth, merged, t.cell, t.sg, 1.0, logger, score);
        REQUIRE(s.ok);
        CHECK((s.twin_decision == "TWINNED") == (fraction > 0));
        if (fraction > 0)
            CHECK(s.r_free_start < s.untwinned_r_free - 0.02);
    }
}