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

#include "RefinementRestraints.h"

#include <algorithm>
#include <cmath>
#include <set>
#include <unordered_map>

#include "gemmi/neighbor.hpp"
#include "gemmi/polyheur.hpp"   // setup_entities
#include "gemmi/riding_h.hpp"   // place_hydrogens_on_all_atoms
#include "gemmi/topo.hpp"

#include "../../common/ParallelFor.h"
#include "MonomerLibrary.h"
#include "RefinementMath.h"

using refine_math::V3;

namespace {

constexpr float CHIRAL_SIGMA = 0.2f;                  // A^3
constexpr float REPULSION_SIGMA = 0.2f;               // A
constexpr float HBOND_CONTACT = 2.6f;                 // A
constexpr float METAL_CONTACT = 1.9f;                 // A, to N, O or S
// B similarity: main chain and the rest, for bonded pairs and pairs bonded to a common atom.
constexpr float B_SIGMA_MAIN_BOND = 1.5f, B_SIGMA_MAIN_ANGLE = 2.0f;
constexpr float B_SIGMA_SIDE_BOND = 3.0f, B_SIGMA_SIDE_ANGLE = 4.5f;

bool IsMainChain(const gemmi::Atom &a) {
    static const std::set<std::string> names = {"N", "CA", "C", "O", "OXT", "P", "OP1", "OP2", "O5'", "C5'",
                                                "C4'", "C3'", "O3'"};
    return names.count(a.name) > 0;
}

bool IsMetal(const gemmi::Element &el) {
    return el.is_metal();
}

// The radius van der Waals contacts are made of, without hydrogens: the dictionary's for the atom's
// energy type, or an element's typical value where the residue is not described.
float VdwRadius(const gemmi::MonLib &monlib, const std::string &resname, const gemmi::Atom &a) {
    auto cc = monlib.monomers.find(resname);
    if (cc != monlib.monomers.end()) {
        auto it = cc->second.find_atom(a.name);
        if (it != cc->second.atoms.end()) {
            auto e = monlib.ener_lib.atoms.find(it->chem_type);
            if (e != monlib.ener_lib.atoms.end() && std::isfinite(e->second.vdw_radius))
                return static_cast<float>(e->second.vdw_radius);
        }
    }
    switch (a.element.elem) {
        case gemmi::El::C: return 1.70f;
        case gemmi::El::N: return 1.55f;
        case gemmi::El::O: return 1.52f;
        case gemmi::El::S: return 1.80f;
        case gemmi::El::P: return 1.80f;
        default: return 1.70f;
    }
}

bool CanHydrogenBond(const gemmi::MonLib &monlib, const std::string &resname, const gemmi::Atom &a) {
    if (a.element != gemmi::El::N && a.element != gemmi::El::O)
        return false;
    auto cc = monlib.monomers.find(resname);
    if (cc != monlib.monomers.end()) {
        auto it = cc->second.find_atom(a.name);
        if (it != cc->second.atoms.end()) {
            auto e = monlib.ener_lib.atoms.find(it->chem_type);
            if (e != monlib.ener_lib.atoms.end())
                return e->second.hb_type == 'D' || e->second.hb_type == 'A' || e->second.hb_type == 'B' ||
                       e->second.hb_type == 'H';
        }
    }
    return true;   // N or O of an undescribed residue
}

V3<double> At(const std::vector<std::array<double, 4>> &xyzb, int i) {
    return {xyzb[i][0], xyzb[i][1], xyzb[i][2]};
}

V3<double> Image(const RestraintRecord &r, V3<double> p) {
    return {r.rot[0] * p.x + r.rot[1] * p.y + r.rot[2] * p.z + r.tran[0],
            r.rot[3] * p.x + r.rot[4] * p.y + r.rot[5] * p.z + r.tran[1],
            r.rot[6] * p.x + r.rot[7] * p.y + r.rot[8] * p.z + r.tran[2]};
}

// The gradient at an image back on the atom it is an image of: rot^T g.
V3<double> FromImage(const RestraintRecord &r, V3<double> g) {
    return {r.rot[0] * g.x + r.rot[3] * g.y + r.rot[6] * g.z,
            r.rot[1] * g.x + r.rot[4] * g.y + r.rot[7] * g.z,
            r.rot[2] * g.x + r.rot[5] * g.y + r.rot[8] * g.z};
}

// A metal ion: a residue of one atom, of a metal element.
bool IsMetalIon(const gemmi::CRA &cra) {
    return cra.residue->atoms.size() == 1 && IsMetal(cra.atom->element);
}

bool IsMetalDonor(const gemmi::Element &el) {
    return el == gemmi::El::O || el == gemmi::El::N || el == gemmi::El::S;
}

// The modification a donor's residue takes when it binds a metal: the proton the bond replaces, as the
// library's own ZN-HISND, ZN-HISNE and ZN-CYS links have it.
std::string MetalDonorModification(const std::string &resname, const std::string &atom) {
    if (resname == "HIS" && atom == "ND1")
        return "DEL-HD1";
    if (resname == "HIS" && atom == "NE2")
        return "DEL-HE2";
    if (resname == "CYS" && atom == "SG")
        return "CYSmod1";
    return "";
}

// The model's LINK / struct_conn records between a metal ion and an O, N or S that no link of the dictionary
// describes (its ZN-CYS, ZN-HISND, ZN-HISNE, FE-CYS, FE-HISND and FE-HISNE are taken as they are): a link of
// one bond is made for each in `monlib`, at the distance of the library's statistics of metal sites
// (EmbeddedMetalDistance) for the two elements and the site's coordination - the number of records between the
// ion and an O, N or S - and the record is pointed at it, so that prepare_topology() restrains it like any other link.
// Only the bond: no angles at the metal. Sites are taken from the records alone, as REFMAC takes links from
// the coordinate file; a contact without a record stays a repulsion (BuildRepulsions, METAL_CONTACT).
// Records to a symmetry mate are left out, as every bond to a symmetry mate is (BuildRefinementTopology).
// Returns the number of links made.
// Distances as surveyed by Harding (2001) Acta Cryst. D57, 401-411 and (2006) Acta Cryst. D62, 678-682,
// and for Mg, Na and K at nucleic acids by Leonarski, D'Ascenzo & Auffinger (2017) Nucleic Acids Res. 45, 987-1004
int AddMetalLinks(gemmi::Structure &st, gemmi::MonLib &monlib) {
    gemmi::Model &model = st.models.at(0);
    struct Site {
        gemmi::Connection *conn;
        gemmi::CRA metal, donor;
        bool metal_first;   // the metal is partner1
    };
    std::vector<Site> sites;
    std::unordered_map<const gemmi::Atom *, int> coordination;
    for (gemmi::Connection &conn : st.connections) {
        gemmi::CRA a = model.find_cra(conn.partner1, true), b = model.find_cra(conn.partner2, true);
        if (!a.atom || !b.atom)
            continue;
        const bool a_metal = IsMetalIon(a), b_metal = IsMetalIon(b);
        if (a_metal == b_metal)
            continue;
        const gemmi::CRA &metal = a_metal ? a : b, &donor = a_metal ? b : a;
        if (!IsMetalDonor(donor.atom->element))
            continue;
        // the site's coordination: every record of the ion, a symmetry mate's or one the dictionary describes too
        coordination[metal.atom]++;
        if (!conn.link_id.empty() || conn.asu == gemmi::Asu::Different ||
            std::get<0>(monlib.match_link(*a.residue, a.atom->name, conn.partner1.altloc, *b.residue, b.atom->name,
                                          conn.partner2.altloc)) != nullptr)
            continue;
        sites.push_back({&conn, metal, donor, a_metal});
    }
    int made = 0;
    for (const Site &site : sites) {
        double distance = 0, sigma = 0;
        if (!EmbeddedMetalDistance(site.metal.atom->element.elem, site.donor.atom->element.elem,
                                   coordination[site.metal.atom], distance, sigma))
            continue;
        const gemmi::Connection &conn = *site.conn;
        gemmi::ChemLink link;
        link.side1.comp = conn.partner1.res_id.name;
        link.side2.comp = conn.partner2.res_id.name;
        (site.metal_first ? link.side2 : link.side1).mod =
            MetalDonorModification(site.donor.residue->name, site.donor.atom->name);
        link.id = "metal-" + std::to_string(made);
        link.name = link.id;
        link.rt.bonds.push_back({gemmi::Restraints::AtomId{1, conn.partner1.atom_name},
                                 gemmi::Restraints::AtomId{2, conn.partner2.atom_name},
                                 gemmi::BondType::Metal, false, distance, sigma, distance, sigma});
        monlib.links.emplace(link.id, link);
        site.conn->link_id = link.id;
        made++;
    }
    return made;
}

}  // namespace

void RestraintSet::Add(RestraintRecord r, int n_atoms_of_r) {
    r.slot = n_slots;
    n_slots += n_atoms_of_r;
    records.push_back(r);
}

void RestraintSet::Append(const RestraintSet &other) {
    const int plane_offset = static_cast<int>(plane_atoms.size());
    for (RestraintRecord r : other.records) {
        r.slot += n_slots;
        if (r.kind == static_cast<int>(RestraintKind::Plane))
            r.atom[0] += plane_offset;
        records.push_back(r);
    }
    plane_atoms.insert(plane_atoms.end(), other.plane_atoms.begin(), other.plane_atoms.end());
    n_slots += other.n_slots;
}

void RestraintSet::Finalize(size_t n_atoms) {
    slot_atom.assign(n_slots, -1);
    for (const RestraintRecord &r : records) {
        if (r.kind == static_cast<int>(RestraintKind::Plane)) {
            for (int k = 0; k < r.atom[1]; k++)
                slot_atom[r.slot + k] = plane_atoms[r.atom[0] + k];
        } else {
            const int n = r.kind == static_cast<int>(RestraintKind::Bond) ||
                                  r.kind == static_cast<int>(RestraintKind::Repulsion) ||
                                  r.kind == static_cast<int>(RestraintKind::BSimilarity) ? 2
                          : r.kind == static_cast<int>(RestraintKind::Angle) ? 3 : 4;
            for (int k = 0; k < n; k++)
                slot_atom[r.slot + k] = r.atom[k];
        }
    }
    slot_start.assign(n_atoms + 1, 0);
    for (int a : slot_atom)
        if (a >= 0)
            slot_start[a + 1]++;
    for (size_t a = 0; a < n_atoms; a++)
        slot_start[a + 1] += slot_start[a];
    slot_index.assign(slot_start[n_atoms], 0);
    std::vector<int> fill(slot_start.begin(), slot_start.end() - 1);
    for (int s = 0; s < n_slots; s++)
        if (slot_atom[s] >= 0)
            slot_index[fill[slot_atom[s]]++] = s;
}

double RestraintValue(const RestraintSet &set, const RestraintRecord &r, const std::vector<std::array<double, 4>> &xyzb) {
    V3<double> g[4];
    switch (static_cast<RestraintKind>(r.kind)) {
        case RestraintKind::Bond:
            return refine_math::Norm(At(xyzb, r.atom[0]) - At(xyzb, r.atom[1]));
        case RestraintKind::Angle: {
            const V3<double> u = At(xyzb, r.atom[0]) - At(xyzb, r.atom[1]), v = At(xyzb, r.atom[2]) - At(xyzb, r.atom[1]);
            return std::atan2(refine_math::Norm(refine_math::Cross(u, v)), refine_math::Dot(u, v)) / refine_math::DEG;
        }
        case RestraintKind::Torsion:
            return refine_math::Dihedral(At(xyzb, r.atom[0]), At(xyzb, r.atom[1]), At(xyzb, r.atom[2]),
                                         At(xyzb, r.atom[3]), g) / refine_math::DEG;
        case RestraintKind::Chiral: {
            const V3<double> c = At(xyzb, r.atom[0]);
            return refine_math::Dot(At(xyzb, r.atom[1]) - c,
                                    refine_math::Cross(At(xyzb, r.atom[2]) - c, At(xyzb, r.atom[3]) - c));
        }
        case RestraintKind::Repulsion:
            return refine_math::Norm(At(xyzb, r.atom[0]) - Image(r, At(xyzb, r.atom[1])));
        case RestraintKind::BSimilarity:
            return xyzb[r.atom[0]][3] - xyzb[r.atom[1]][3];
        case RestraintKind::Plane:
            break;
    }
    (void) set;
    return NAN;
}

double EvaluateRestraints(const RestraintSet &set, const std::vector<std::array<double, 4>> &xyzb,
                          std::vector<std::array<double, 4>> *grad) {
    double sum = 0;
    std::vector<V3<double>> px, pg;
    for (const RestraintRecord &r : set.records) {
        V3<double> g[4] = {};
        int n = 0;
        double e = 0;
        const double ideal = r.ideal, sigma = r.sigma;
        switch (static_cast<RestraintKind>(r.kind)) {
            case RestraintKind::Bond:
                e = refine_math::BondTerm(At(xyzb, r.atom[0]), At(xyzb, r.atom[1]), ideal, sigma, g);
                n = 2;
                break;
            case RestraintKind::Angle:
                e = refine_math::AngleTerm(At(xyzb, r.atom[0]), At(xyzb, r.atom[1]), At(xyzb, r.atom[2]), ideal, sigma, g);
                n = 3;
                break;
            case RestraintKind::Torsion:
                e = refine_math::TorsionTerm(At(xyzb, r.atom[0]), At(xyzb, r.atom[1]), At(xyzb, r.atom[2]),
                                             At(xyzb, r.atom[3]), ideal, sigma, r.period, g);
                n = 4;
                break;
            case RestraintKind::Chiral:
                e = refine_math::ChiralTerm(At(xyzb, r.atom[0]), At(xyzb, r.atom[1]), At(xyzb, r.atom[2]),
                                            At(xyzb, r.atom[3]), ideal, sigma, g);
                n = 4;
                break;
            case RestraintKind::Repulsion:
                e = refine_math::RepulsionTerm(At(xyzb, r.atom[0]), Image(r, At(xyzb, r.atom[1])), ideal, sigma, g);
                g[1] = FromImage(r, g[1]);
                n = 2;
                break;
            case RestraintKind::Plane: {
                const int m = r.atom[1];
                px.resize(m);
                pg.resize(m);
                for (int k = 0; k < m; k++)
                    px[k] = At(xyzb, set.plane_atoms[r.atom[0] + k]);
                e = refine_math::PlaneTerm(px.data(), m, sigma, pg.data());
                if (grad != nullptr)
                    for (int k = 0; k < m; k++) {
                        std::array<double, 4> &d = (*grad)[set.plane_atoms[r.atom[0] + k]];
                        d[0] += pg[k].x;
                        d[1] += pg[k].y;
                        d[2] += pg[k].z;
                    }
                break;
            }
            case RestraintKind::BSimilarity: {
                const double rr = (xyzb[r.atom[0]][3] - xyzb[r.atom[1]][3]) / sigma;
                e = rr * rr;
                if (grad != nullptr) {
                    (*grad)[r.atom[0]][3] += 2 * rr / sigma;
                    (*grad)[r.atom[1]][3] -= 2 * rr / sigma;
                }
                break;
            }
        }
        sum += e;
        if (grad != nullptr)
            for (int k = 0; k < n; k++) {
                std::array<double, 4> &d = (*grad)[r.atom[k]];
                d[0] += g[k].x;
                d[1] += g[k].y;
                d[2] += g[k].z;
            }
    }
    return sum;
}

std::vector<double> BRestraintCurvature(const RestraintSet &set, const std::vector<std::array<double, 4>> &xyzb) {
    std::vector<double> h(xyzb.size(), 0.0);
    for (const RestraintRecord &r : set.records)
        if (r.kind == static_cast<int>(RestraintKind::BSimilarity)) {
            h[r.atom[0]] += 2.0 / (r.sigma * r.sigma);
            h[r.atom[1]] += 2.0 / (r.sigma * r.sigma);
        }
    return h;
}

RefinementTopology BuildRefinementTopology(gemmi::Structure &st, bool hydrogens) {
    RefinementTopology out;
    gemmi::Model &model = st.models.at(0);
    out.monlib = std::make_shared<gemmi::MonLib>(EmbeddedMonomerLibrary());
    gemmi::MonLib &monlib = *out.monlib;
    gemmi::setup_entities(st);
    out.metal_links = AddMetalLinks(st, monlib);
    gemmi::Logger quiet;   // missing atoms and unknown residues are what the fixed residues report
    quiet.threshold = 0;
    out.topo = gemmi::prepare_topology(st, monlib, 0,
                                       hydrogens ? gemmi::HydrogenChange::ReAddButWater : gemmi::HydrogenChange::NoChange,
                                       /*reorder=*/false, quiet, /*ignore_unknown_links=*/true, /*use_cispeps=*/true);
    gemmi::Topo *topo = out.topo.get();

    // the atoms as they are now, hydrogens added
    std::unordered_map<const gemmi::Atom *, int> index;
    std::vector<const gemmi::Residue *> residue_of;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            for (const gemmi::Atom &a : r.atoms) {
                index[&a] = static_cast<int>(residue_of.size());
                residue_of.push_back(&r);
            }
    const size_t n = residue_of.size();
    out.refinable.assign(n, 0);
    out.vdw_radius.assign(n, 1.7f);
    out.hbond.assign(n, 0);
    out.metal.assign(n, 0);
    out.altloc.assign(n, 0);
    out.element.assign(n, gemmi::El::X);
    out.rider_parent.assign(n, -1);
    for (const gemmi::Topo::Bond &b : topo->bonds)
        for (int k = 0; k < 2; k++)
            if (b.atoms[k]->is_hydrogen() && !b.atoms[1 - k]->is_hydrogen())
                out.rider_parent[index.at(b.atoms[k])] = index.at(b.atoms[1 - k]);
    // A hydrogen gemmi could not place, or with no parent, scatters nothing.
    for (gemmi::Chain &ch : model.chains)
        for (gemmi::Residue &r : ch.residues)
            for (gemmi::Atom &a : r.atoms)
                if (a.is_hydrogen() && (out.rider_parent[index.at(&a)] < 0 || !std::isfinite(a.pos.x))) {
                    a.occ = 0;
                    out.rider_parent[index.at(&a)] = -1;
                    if (!std::isfinite(a.pos.x))
                        a.pos = r.atoms.front().pos;
                }

    // A residue is refined where the dictionary describes it and every one of its atoms.
    std::set<std::string> fixed_names;
    for (const gemmi::Topo::ChainInfo &ci : topo->chain_infos)
        for (const gemmi::Topo::ResInfo &ri : ci.res_infos) {
            const gemmi::Residue &res = *ri.res;
            bool known = ri.orig_chemcomp != nullptr && InEmbeddedMonomerLibrary(res.name);
            for (const gemmi::Atom &a : res.atoms)
                known = known && ri.get_final_chemcomp(a.altloc).has_atom(a.name);
            if (known)
                out.restrained_residues++;
            else {
                out.fixed_residues++;
                fixed_names.insert(res.name);
            }
            for (const gemmi::Atom &a : res.atoms) {
                const auto it = index.find(&a);
                if (it == index.end())
                    continue;
                const int i = it->second;
                out.refinable[i] = known && a.occ > 0 && !a.is_hydrogen();
                out.vdw_radius[i] = VdwRadius(monlib, res.name, a);
                out.hbond[i] = res.is_water() || CanHydrogenBond(monlib, res.name, a);
                out.metal[i] = res.atoms.size() == 1 && IsMetal(a.element);
                out.altloc[i] = a.altloc;
                out.element[i] = a.element.elem;
            }
        }
    out.fixed_names.assign(fixed_names.begin(), fixed_names.end());

    auto idx = [&](const gemmi::Atom *a) { return index.at(a); };
    auto has_h = [](std::initializer_list<const gemmi::Atom *> atoms) {
        for (const gemmi::Atom *a : atoms)
            if (a->is_hydrogen())
                return true;
        return false;
    };
    auto any_refinable = [&](std::initializer_list<int> atoms) {
        for (int a : atoms)
            if (out.refinable[a])
                return true;
        return false;
    };
    std::vector<std::set<int>> excluded(n);
    std::vector<std::vector<int>> bonded(n);
    for (size_t i = 0; i < n; i++)
        excluded[i].insert(static_cast<int>(i));

    for (const gemmi::Topo::Bond &b : topo->bonds) {
        if (b.asu == gemmi::Asu::Different || has_h({b.atoms[0], b.atoms[1]}))
            continue;
        const int i = idx(b.atoms[0]), j = idx(b.atoms[1]);
        excluded[i].insert(j);
        excluded[j].insert(i);
        bonded[i].push_back(j);
        bonded[j].push_back(i);
        if (!any_refinable({i, j}) || !(b.restr->esd > 0))
            continue;
        RestraintRecord r;
        r.kind = static_cast<int>(RestraintKind::Bond);
        r.atom[0] = i;
        r.atom[1] = j;
        r.ideal = static_cast<float>(b.restr->value);
        r.sigma = static_cast<float>(b.restr->esd);
        if (out.metal[i] || out.metal[j])
            out.metal_bonds.Add(r, 2);   // a metal's coordination, kept apart from the covalent bonds
        else {
            out.geometry.Add(r, 2);
            out.bonds++;
        }
        RestraintRecord rb;
        rb.kind = static_cast<int>(RestraintKind::BSimilarity);
        rb.atom[0] = i;
        rb.atom[1] = j;
        const bool main = IsMainChain(*b.atoms[0]) && IsMainChain(*b.atoms[1]);
        rb.sigma = main ? B_SIGMA_MAIN_BOND : B_SIGMA_SIDE_BOND;
        out.adp.Add(rb, 2);
    }
    for (const gemmi::Topo::Angle &an : topo->angles) {
        if (has_h({an.atoms[0], an.atoms[1], an.atoms[2]}))
            continue;
        const int i = idx(an.atoms[0]), j = idx(an.atoms[1]), k = idx(an.atoms[2]);
        excluded[i].insert(k);
        excluded[k].insert(i);
        if (!any_refinable({i, j, k}) || !(an.restr->esd > 0))
            continue;
        RestraintRecord r;
        r.kind = static_cast<int>(RestraintKind::Angle);
        r.atom[0] = i;
        r.atom[1] = j;
        r.atom[2] = k;
        r.ideal = static_cast<float>(an.restr->value);
        r.sigma = static_cast<float>(an.restr->esd);
        out.geometry.Add(r, 3);
        out.angles++;
        RestraintRecord rb;
        rb.kind = static_cast<int>(RestraintKind::BSimilarity);
        rb.atom[0] = i;
        rb.atom[1] = k;
        const bool main = IsMainChain(*an.atoms[0]) && IsMainChain(*an.atoms[2]);
        rb.sigma = main ? B_SIGMA_MAIN_ANGLE : B_SIGMA_SIDE_ANGLE;
        out.adp.Add(rb, 2);
    }
    for (const gemmi::Topo::Torsion &t : topo->torsions) {
        const std::string &label = t.restr->label;
        if (has_h({t.atoms[0], t.atoms[1], t.atoms[2], t.atoms[3]}))
            continue;
        if (!(t.restr->esd > 0) || label.rfind("phi", 0) == 0 || label.rfind("psi", 0) == 0 ||
            label.rfind("C2e", 0) == 0 || label.rfind("C3e", 0) == 0)
            continue;
        const int a = idx(t.atoms[0]), b = idx(t.atoms[1]), c = idx(t.atoms[2]), d = idx(t.atoms[3]);
        if (!any_refinable({a, b, c, d}))
            continue;
        RestraintRecord r;
        r.kind = static_cast<int>(RestraintKind::Torsion);
        r.atom[0] = a;
        r.atom[1] = b;
        r.atom[2] = c;
        r.atom[3] = d;
        r.ideal = static_cast<float>(t.restr->value);
        r.sigma = static_cast<float>(t.restr->esd);
        r.period = t.restr->period;
        out.geometry.Add(r, 4);
        out.torsions++;
    }
    for (const gemmi::Topo::Chirality &ch : topo->chirs) {
        if (ch.restr->sign == gemmi::ChiralityType::Both || has_h({ch.atoms[0], ch.atoms[1], ch.atoms[2], ch.atoms[3]}))
            continue;
        const int c = idx(ch.atoms[0]), a1 = idx(ch.atoms[1]), a2 = idx(ch.atoms[2]), a3 = idx(ch.atoms[3]);
        if (!any_refinable({c, a1, a2, a3}))
            continue;
        const double volume = topo->ideal_chiral_abs_volume(ch);
        RestraintRecord r;
        r.kind = static_cast<int>(RestraintKind::Chiral);
        r.atom[0] = c;
        r.atom[1] = a1;
        r.atom[2] = a2;
        r.atom[3] = a3;
        r.ideal = static_cast<float>(ch.restr->sign == gemmi::ChiralityType::Negative ? -volume : volume);
        r.sigma = CHIRAL_SIGMA;
        out.geometry.Add(r, 4);
        out.chirals++;
    }
    for (const gemmi::Topo::Plane &p_all : topo->planes) {
        gemmi::Topo::Plane p = p_all;
        std::erase_if(p.atoms, [](const gemmi::Atom *a) { return a->is_hydrogen(); });
        if (p.atoms.size() < 4 || !(p.restr->esd > 0))
            continue;
        bool any = false;
        for (const gemmi::Atom *a : p.atoms)
            any = any || out.refinable[idx(a)];
        if (!any)
            continue;
        RestraintRecord r;
        r.kind = static_cast<int>(RestraintKind::Plane);
        r.atom[0] = static_cast<int>(out.geometry.plane_atoms.size());
        r.atom[1] = static_cast<int>(p.atoms.size());
        r.sigma = static_cast<float>(p.restr->esd);
        for (const gemmi::Atom *a : p.atoms)
            out.geometry.plane_atoms.push_back(idx(a));
        out.geometry.Add(r, r.atom[1]);
        out.planes++;
    }
    // and those three bonds away, whose distance the torsions and planes govern
    std::vector<std::set<int>> within3 = excluded;
    for (size_t i = 0; i < n; i++)
        for (int j : excluded[i])
            for (int k : bonded[j])
                within3[i].insert(k);
    out.excluded.resize(n);
    for (size_t i = 0; i < n; i++)
        out.excluded[i].assign(within3[i].begin(), within3[i].end());
    out.geometry.Finalize(n);
    out.metal_bonds.Finalize(n);
    out.adp.Finalize(n);
    return out;
}

RestraintSet BuildRepulsions(const RefinementTopology &topo, gemmi::Model &model, const gemmi::UnitCell &cell_in,
                             const gemmi::SpaceGroup &sg, double margin, size_t nthreads) {
    gemmi::UnitCell cell = cell_in;
    cell.set_cell_images_from_spacegroup(&sg);
    float max_radius = 0;
    for (float r : topo.vdw_radius)
        max_radius = std::max(max_radius, r);
    const double reach = 2 * max_radius + margin;
    gemmi::NeighborSearch ns(model, cell, std::max(reach, 5.0));
    ns.populate(false);

    std::vector<std::vector<int>> first(model.chains.size());
    std::vector<gemmi::Position> pos;
    for (size_t c = 0; c < model.chains.size(); c++)
        for (const gemmi::Residue &r : model.chains[c].residues) {
            first[c].push_back(static_cast<int>(pos.size()));
            for (const gemmi::Atom &a : r.atoms)
                pos.push_back(a.pos);
        }

    // The atoms in blocks searched on the threads, each block's pairs listed in atom order; the blocks are
    // joined in order below, so the list is the same for any number of threads.
    constexpr int BLOCKS = 64;
    std::vector<std::vector<RestraintRecord>> found(BLOCKS);
    ParallelFor(BLOCKS, nthreads, [&](int block) {
        const size_t lo = pos.size() * block / BLOCKS, hi = pos.size() * (block + 1) / BLOCKS;
        for (size_t i = lo; i < hi; i++) {
            if (topo.rider_parent[i] >= 0 || topo.element[i] == gemmi::El::H)
                continue;
            const gemmi::Fractional fi = cell.fractionalize(pos[i]);
            ns.for_each(pos[i], '\0', reach, [&](gemmi::NeighborSearch::Mark &m, double) {
                const int j = first[m.chain_idx][m.residue_idx] + m.atom_idx;
                if (topo.rider_parent[j] >= 0 || topo.element[j] == gemmi::El::H)
                    return;
                if (!topo.refinable[i] && !topo.refinable[j])
                    return;
                if (m.image_idx == 0 && (static_cast<int>(i) >= j ||
                                         std::binary_search(topo.excluded[i].begin(), topo.excluded[i].end(), j)))
                    return;
                if (m.image_idx != 0 && static_cast<int>(i) > j)
                    return;
                if (topo.altloc[i] && topo.altloc[j] && topo.altloc[i] != topo.altloc[j])
                    return;
                // The image of j nearest to i, as a Cartesian x -> R x + t.
                const gemmi::FTransform ft = ns.get_image_transformation(m.image_idx);
                const gemmi::Fractional fj = ft.apply(cell.fractionalize(pos[j]));
                const gemmi::Vec3 shift(std::round(fi.x - fj.x), std::round(fi.y - fj.y), std::round(fi.z - fj.z));
                if (m.image_idx == 0 && (shift.x != 0 || shift.y != 0 || shift.z != 0))
                    return;   // the same copy's own lattice translate: the identity term covers the contact
                const gemmi::Mat33 rot = cell.orth.mat.multiply(ft.mat).multiply(cell.frac.mat);
                const gemmi::Vec3 tran = cell.orth.mat.multiply(ft.vec + shift);
                const gemmi::Position image(rot.multiply(pos[j]) + tran);
                const double d = pos[i].dist(image);

                float d0 = topo.vdw_radius[i] + topo.vdw_radius[j];
                if (topo.metal[i] || topo.metal[j]) {
                    const gemmi::El other = topo.element[topo.metal[i] ? j : static_cast<int>(i)];
                    if (other == gemmi::El::N || other == gemmi::El::O || other == gemmi::El::S)
                        d0 = METAL_CONTACT;
                } else if (topo.hbond[i] && topo.hbond[j])
                    d0 = HBOND_CONTACT;
                if (!(d < d0 + margin))
                    return;
                RestraintRecord r;
                r.kind = static_cast<int>(RestraintKind::Repulsion);
                r.atom[0] = static_cast<int>(i);
                r.atom[1] = j;
                r.ideal = d0;
                r.sigma = REPULSION_SIGMA;
                for (int a = 0; a < 3; a++) {
                    for (int b = 0; b < 3; b++)
                        r.rot[3 * a + b] = static_cast<float>(rot[a][b]);
                    r.tran[a] = static_cast<float>(tran.at(a));
                }
                found[block].push_back(r);
            });
        }
    });
    RestraintSet out;
    for (const std::vector<RestraintRecord> &records : found)
        for (const RestraintRecord &r : records)
            out.Add(r, 2);
    out.Finalize(pos.size());
    return out;
}

void PlaceRidingHydrogens(RefinementTopology &topo) {
    gemmi::place_hydrogens_on_all_atoms(*topo.topo);
}

GeometryStatistics MeasureGeometry(const RestraintSet &geometry, const RestraintSet &repulsions,
                                   const std::vector<std::array<double, 4>> &xyzb, const RestraintSet *metal_bonds) {
    GeometryStatistics s;
    if (metal_bonds != nullptr) {
        double sm = 0, zm = 0;
        for (const RestraintRecord &r : metal_bonds->records) {
            const double d = RestraintValue(*metal_bonds, r, xyzb) - r.ideal;
            sm += d * d;
            zm += d * d / (double(r.sigma) * r.sigma);
            s.metal_bonds++;
        }
        s.rmsd_metal = s.metal_bonds > 0 ? std::sqrt(sm / s.metal_bonds) : 0;
        s.rmsz_metal = s.metal_bonds > 0 ? std::sqrt(zm / s.metal_bonds) : 0;
    }
    double sb = 0, sa = 0, sc = 0, zb = 0, za = 0;
    int nc = 0;
    for (const RestraintRecord &r : geometry.records) {
        const RestraintKind kind = static_cast<RestraintKind>(r.kind);
        if (kind == RestraintKind::Bond) {
            const double d = RestraintValue(geometry, r, xyzb) - r.ideal;
            sb += d * d;
            zb += d * d / (double(r.sigma) * r.sigma);
            s.bonds++;
        } else if (kind == RestraintKind::Angle) {
            const double d = RestraintValue(geometry, r, xyzb) - r.ideal;
            sa += d * d;
            za += d * d / (double(r.sigma) * r.sigma);
            s.angles++;
        } else if (kind == RestraintKind::Chiral) {
            const double d = RestraintValue(geometry, r, xyzb) - r.ideal;
            sc += d * d;
            nc++;
        }
    }
    for (const RestraintRecord &r : repulsions.records)
        if (RestraintValue(repulsions, r, xyzb) < r.ideal - 0.4)
            s.clashes++;
    s.rmsd_bonds = s.bonds > 0 ? std::sqrt(sb / s.bonds) : 0;
    s.rmsd_angles = s.angles > 0 ? std::sqrt(sa / s.angles) : 0;
    s.rmsd_chirals = nc > 0 ? std::sqrt(sc / nc) : 0;
    s.rmsz_bonds = s.bonds > 0 ? std::sqrt(zb / s.bonds) : 0;
    s.rmsz_angles = s.angles > 0 ? std::sqrt(za / s.angles) : 0;
    return s;
}