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

#include "SolventSites.h"

#include <algorithm>
#include <cmath>
#include <map>
#include <mutex>
#include <set>
#include <tuple>

#include <spdlog/fmt/fmt.h>

#include "gemmi/calculate.hpp"   // find_best_plane
#include "gemmi/fprime.hpp"      // cromer_liberman
#include "gemmi/neighbor.hpp"
#include "gemmi/resinfo.hpp"

#include "../../common/ParallelFor.h"
#include "ModelGrid.h"             // IonFormFactor
#include "MonomerLibrary.h"
#include "RefinementRestraints.h"   // AtomsNear, SiteSymmetry

namespace {

constexpr double PI = 3.14159265358979;
constexpr double SEARCH_RADIUS = 6.0;        // A: ion pairs are listed out to it, as Cat_Wiz lists them
constexpr double SHELL_RADIUS = 3.4;         // A: the N and O neighbours listed and typed
// A: the ligands bond valences are summed over, following Echols et al. (2014) Acta Cryst. D70, 1104-1114
constexpr double BVS_RADIUS = 3.5;
constexpr double LIKE_CHARGE_DISTANCE = 3.5; // A: like charges closer than this repel
// A: no negative group this close to a halide, following Morshed et al. (2015) Acta Cryst. D71, 1147-1158
constexpr double HALIDE_NEGATIVE_DISTANCE = 3.2;
constexpr double BV_B = 0.37;                // A, the bond-valence softness
// A band's edges are soft by this much (A); a neighbour's coordinate error widens the band beyond it.
constexpr double BAND_EDGE = 0.05;
// A contact the species does not make (closer than its shell, in its exclusion zone, a donor at a cation or an
// acceptor at a halide, a like charge), in the score's log units, per contact.
constexpr double BAD_CONTACT = 2.0;
// A first-shell ligand fewer than a Mg, Na, Ca or Cl holds: a water the model may simply not have, so it costs
// little. Incomplete shells are common; they are no evidence for the ion, and not much against it - and a halide
// binds at hydrophobic surfaces too, with few or no donors, where its electrons and its anomalous signal, not its
// shell, tell it from a water.
// Following Dauter & Dauter (2001) Structure 9, R21-R26
constexpr double MISSING_LIGAND = 0.5;
// The electron count's calibration on the waters: on those within SEARCH_RADIUS of the site where there are at
// least this many, or on all of the model's; its spread (rms of the log of the waters' calibrations) at least this.
// The density the model's other atoms put at a point is counted out to TAIL_RADIUS.
constexpr int CALIBRATION_NEIGHBOURS = 3;
constexpr double TAIL_RADIUS = 4.5;   // A
constexpr double ELECTRON_SPREAD_MIN = 0.15;
// A divalent ion on a purine N7: within 1.0 A of the base's plane and 3.5-3.9 A from its O6 or N6, following
// Leonarski, D'Ascenzo & Auffinger (2017) Nucleic Acids Res. 45, 987-1004
constexpr double N7_PLANE_MAX = 1.0, N7_O6N6_MIN = 3.5, N7_O6N6_MAX = 3.9;

// R0 (A) of the bond-valence parameters for Mg, Na, K, Ca to O and to N, b = 0.37 A, as tabulated by Zheng et al.
// (2017) Acta Cryst. D73, 316-325, Table 2.
// Following Brese & O'Keeffe (1991) Acta Cryst. B47, 192-197
constexpr double R0_O[BVS_CATIONS] = {1.69, 1.80, 2.13, 1.96};
constexpr double R0_N[BVS_CATIONS] = {1.85, 1.93, 2.26, 2.14};

// What a species' shell looks like: its first shell of N and O [lo, hi) and how many it holds, the distance to O
// and to N with its spread where the papers give one (Mg only: elsewhere the band is all there is, esd 0), the N
// and O it cannot have nearer than `too_close`, an exclusion zone [exclude_lo, exclude_hi) it keeps empty (Mg
// only), its charge and electrons, and its row in SolventSite::bvs (-1 none).
// Coordination bands following Naleem et al. (2026) Nucleic Acids Res. 54, gkag120 (Cat_Wiz)
// Distances and exclusion zone following Leonarski et al. (2017) Nucleic Acids Res. 45, 987-1004, (2019) RNA 25,
// 173-192 and (2025) Nucleic Acids Res. 53, gkae1148
struct SpeciesRule {
    double lo, hi;
    double ideal, esd, ideal_n, esd_n;
    double too_close;
    double exclude_lo, exclude_hi;
    double cn, cn_sigma;
    int charge;
    int electrons;
    int bvs;
};

const SpeciesRule RULES[SOLVENT_SPECIES] = {
        // water: 2.6-3.2 A, at most four (two donated, two accepted), at least one
        {2.6, 3.2, 0, 0, 0, 0, 2.6, 0, 0, 4, 1.0, 0, 10, -1},
        // Mg2+: six at 1.9-2.3 A, Mg-O 2.07 +- 0.04 A, Mg-N 2.19 +- 0.08 A, no N or O at 2.3-3.4 A
        {1.9, 2.3, 2.07, 0.04, 2.19, 0.08, 1.9, 2.3, 3.4, 6, 0.7, 2, 10, 0},
        // Na+: five or six at 2.3-2.6 A
        {2.3, 2.6, 0, 0, 0, 0, 2.3, 0, 0, 5.5, 1.0, 1, 10, 1},
        // K+: six to ten at 2.6-3.2 A
        {2.6, 3.2, 0, 0, 0, 0, 2.6, 0, 0, 8, 1.5, 1, 18, 2},
        // Ca2+: six to eight at 2.3-2.6 A
        {2.3, 2.6, 0, 0, 0, 0, 2.3, 0, 0, 7, 1.0, 2, 18, 3},
        // Cl-: three to five hydrogen-bond donors at 3.0-3.4 A, almost nothing nearer than 2.9 A
        // Following Carugo (2014) BMC Struct. Biol. 14, 19
        {3.0, 3.4, 0, 0, 0, 0, 2.9, 0, 0, 4, 1.0, -1, 18, -1},
        // NH4+: four acceptors, at water's distances
        {2.6, 3.2, 0, 0, 0, 0, 2.6, 0, 0, 4, 1.0, 1, 10, -1},
};

bool IsWater(const gemmi::Residue &r) {
    return r.is_water();
}

// The element, charge and riding hydrogens each species scatters as: a water as O with its two H, NH4+ as N
// with four.
struct Scatterer {
    gemmi::El el;
    signed char charge;
    int hydrogens;
};
const Scatterer SCATTERERS[SOLVENT_SPECIES] = {{gemmi::El::O, 0, 2},  {gemmi::El::Mg, 2, 0}, {gemmi::El::Na, 1, 0},
                                               {gemmi::El::K, 1, 0},  {gemmi::El::Ca, 2, 0}, {gemmi::El::Cl, -1, 0},
                                               {gemmi::El::N, 0, 4}};

// The height at its centre of the density of an atom of form factor f(s) and B, as a map to d_min shows it:
// 4 pi int_0^(1/d_min) f(s) exp(-B s^2 / 4) s^2 ds (s = 1/d), by Simpson's rule.
// At distance r from the centre: 4 pi int f(s) exp(-B s^2 / 4) s^2 sin(2 pi s r) / (2 pi s r) ds.
template <class F>
double PeakHeight(F &&f, double b, double d_min, double r = 0) {
    constexpr int N = 64;
    const double s_max = 1.0 / d_min, h = s_max / N;
    double sum = 0;
    for (int k = 0; k <= N; k++) {
        const double s = k * h;
        const double w = (k == 0 || k == N) ? 1 : (k % 2 == 1 ? 4 : 2);
        const double x = 2 * PI * s * r;
        sum += w * f(s) * std::exp(-b * s * s / 4) * s * s * (x > 1e-9 ? std::sin(x) / x : 1.0);
    }
    return 4 * PI * sum * h / 3;
}

double ScattererHeight(const Scatterer &sc, double b, double d_min, double r = 0) {
    const auto &f = IonFormFactor<double>(sc.el, sc.charge);
    const auto &fh = IonFormFactor<double>(gemmi::El::H, 0);
    return PeakHeight([&](double s) { return f.calculate_sf(s * s / 4) + sc.hydrogens * fh.calculate_sf(s * s / 4); },
                      b, d_min, r);
}

// How an atom of the model scatters for the density it puts at its neighbours: as modelled, a water with its
// hydrogens.
Scatterer ScattererOf(const gemmi::Residue &r, const gemmi::Atom &a) {
    return {a.element.elem, a.charge, IsWater(r) ? 2 : 0};
}

// The density of an atom at distance r from it (ScattererHeight), tabulated per scatterer and B (to the nearest
// A^2) at steps of 0.05 A out to TAIL_RADIUS, read by linear interpolation.
class RadialDensity {
public:
    explicit RadialDensity(double d_min) : d_min_(d_min) {}
    double operator()(const Scatterer &sc, double b, double r) { return At(Table(sc, b), r); }
    // The tables of these scatterers at these B made up front, on `nthreads` threads: Table() then only reads them.
    void Prepare(const std::vector<std::pair<Scatterer, double>> &wanted, size_t nthreads) {
        std::map<Key, std::vector<double>> made;
        for (const auto &[sc, b] : wanted)
            made[KeyOf(sc, b)];
        std::vector<std::pair<const Key, std::vector<double>> *> todo;
        for (auto &entry : made)
            if (!tables_.count(entry.first))
                todo.push_back(&entry);
        ParallelFor(static_cast<int>(todo.size()), nthreads, [&](int k) { todo[k]->second = Make(todo[k]->first); });
        for (auto *entry : todo)
            tables_.emplace(entry->first, std::move(entry->second));
    }
    const std::vector<double> &Table(const Scatterer &sc, double b) {
        const std::lock_guard lock(m_);   // the sites are analysed on several threads; a table, once made, stays
        const Key key = KeyOf(sc, b);
        auto it = tables_.find(key);
        if (it == tables_.end())
            it = tables_.emplace(key, Make(key)).first;
        return it->second;
    }
    static double At(const std::vector<double> &t, double r) {
        if (!(r < TAIL_RADIUS))
            return 0;
        const double x = r / STEP;
        const size_t k = static_cast<size_t>(x);
        return t[k] + (x - k) * (t[k + 1] - t[k]);
    }

private:
    using Key = std::tuple<int, int, int, int>;   // element, charge, hydrogens, B to the nearest A^2
    static Key KeyOf(const Scatterer &sc, double b) {
        return {static_cast<int>(sc.el), static_cast<int>(sc.charge), sc.hydrogens,
                static_cast<int>(std::lround(std::clamp(b, 0.0, 500.0)))};
    }
    std::vector<double> Make(const Key &key) const {
        const Scatterer sc{static_cast<gemmi::El>(std::get<0>(key)), static_cast<signed char>(std::get<1>(key)),
                           std::get<2>(key)};
        std::vector<double> t(static_cast<size_t>(TAIL_RADIUS / STEP) + 2);
        for (size_t k = 0; k < t.size(); k++)
            t[k] = ScattererHeight(sc, std::get<3>(key), d_min_, k * STEP);
        return t;
    }
    static constexpr double STEP = 0.05;
    double d_min_;
    std::mutex m_;
    std::map<Key, std::vector<double>> tables_;
};

// An atom near a site, as it scatters (for the density it puts there).
struct Scattering {
    gemmi::Position pos;
    double occ;
    const std::vector<double> *density;   // RadialDensity::Table() of its scatterer at its B
};

// The site's own density fitted to the map near it, the model's other atoms (`others`) held: 2mFo-DFc, on the model's
// scale (`k` per model unit), at points 0.3 A apart within FIT_RADIUS of `p0`, against `a` times a water (O and two
// H) at B `b` centred at p plus the others' density there, least squares in a; p is p0, or with `search` the best of
// the points 0.1 A apart within SOLVENT_PEAK_REACH of it. The site's height thus counted apart from its neighbours',
// which merge with it at moderate resolution: a is the electrons it shows, in waters.
struct LocalFit {
    gemmi::Position p;
    double a = NAN;
};
constexpr double FIT_RADIUS = 1.0;   // A

LocalFit FitSite(const gemmi::Grid<float> &map, double k, const gemmi::Position &p0, double b,
                 const std::vector<Scattering> &others, RadialDensity &radial, bool search) {
    std::vector<gemmi::Position> points;
    std::vector<double> obs;
    constexpr double STEP = 0.3;
    const double reach = FIT_RADIUS + (search ? SOLVENT_PEAK_REACH : 0.0);
    const int n = static_cast<int>(std::ceil(reach / STEP));
    for (int i = -n; i <= n; i++)
        for (int j = -n; j <= n; j++)
            for (int l = -n; l <= n; l++) {
                const gemmi::Position x(p0.x + i * STEP, p0.y + j * STEP, p0.z + l * STEP);
                if (x.dist(p0) > reach)
                    continue;
                double tail = 0;
                for (const Scattering &o : others)
                    tail += o.occ * RadialDensity::At(*o.density, x.dist(o.pos));
                points.push_back(x);
                obs.push_back(map.interpolate_value(x) / k - tail);
            }
    const std::vector<double> &water = radial.Table(SCATTERERS[0], b);
    auto fit = [&](const gemmi::Position &p, double &a) {
        double num = 0, den = 0, oo = 0;
        for (size_t m = 0; m < points.size(); m++)
            if (const double d = points[m].dist(p); d <= FIT_RADIUS) {
                const double shape = RadialDensity::At(water, d);
                num += obs[m] * shape;
                den += shape * shape;
                oo += obs[m] * obs[m];
            }
        a = den > 0 ? num / den : NAN;
        return den > 0 ? oo - num * num / den : INFINITY;   // the residual at the best a
    };
    LocalFit out;
    out.p = p0;
    double best = fit(p0, out.a);
    // the points 0.2 A apart within reach, then 0.1 A apart around the best of them
    for (double step : {0.2, 0.1}) {
        if (!search)
            break;
        const gemmi::Position centre = out.p;
        const int m = step > 0.15 ? static_cast<int>(std::lround(SOLVENT_PEAK_REACH / step)) : 1;
        for (int i = -m; i <= m; i++)
            for (int j = -m; j <= m; j++)
                for (int l = -m; l <= m; l++) {
                    const gemmi::Position p(centre.x + i * step, centre.y + j * step, centre.z + l * step);
                    double a = 0;
                    if (p.dist(p0) > SOLVENT_PEAK_REACH)
                        continue;
                    if (const double r = fit(p, a); r < best && a > 0) {
                        best = r;
                        out = {p, a};
                    }
                }
    }
    return out;
}

// The median and the rms spread (1.4826 times the median absolute deviation) of `v`.
std::pair<double, double> MedianSpread(std::vector<double> v) {
    if (v.empty())
        return {NAN, NAN};
    std::sort(v.begin(), v.end());
    const double m = v[v.size() / 2];
    std::vector<double> dev;
    for (double x : v)
        dev.push_back(std::fabs(x - m));
    std::sort(dev.begin(), dev.end());
    return {m, 1.4826 * dev[dev.size() / 2]};
}

bool IsHalogen(gemmi::El el) {
    return el == gemmi::El::F || el == gemmi::El::Cl || el == gemmi::El::Br || el == gemmi::El::I;
}

bool Polar(gemmi::El el) {
    return el == gemmi::El::N || el == gemmi::El::O;
}

// The species a residue name labels, -1 for one not scored here.
int LabelSpecies(const std::string &name) {
    if (name == "HOH" || name == "WAT" || name == "DOD" || name == "H2O")
        return static_cast<int>(SolventSpecies::Water);
    for (int s = 1; s < SOLVENT_SPECIES; s++)
        if (name == SolventSpeciesName(static_cast<SolventSpecies>(s)))
            return s;
    return -1;
}

// The one heavy atom of a single-atom ion, or of NH4; null for any other residue.
const gemmi::Atom *IonAtom(const gemmi::Residue &r) {
    const gemmi::Atom *heavy = nullptr;
    for (const gemmi::Atom &a : r.atoms)
        if (!a.is_hydrogen()) {
            if (heavy != nullptr && heavy->altloc == a.altloc)
                return nullptr;
            heavy = &a;
        }
    if (heavy == nullptr)
        return nullptr;
    if (r.name == "NH4" || heavy->element.is_metal() || IsHalogen(heavy->element.elem))
        return heavy;
    return nullptr;
}

// +1 a cation, -1 an anion, 0 a water or not a solvent site.
int IonCharge(const gemmi::Residue &r) {
    if (IsWater(r))
        return 0;
    const gemmi::Atom *a = IonAtom(r);
    if (a == nullptr)
        return 0;
    return IsHalogen(a->element.elem) ? -1 : 1;
}

// Probability that a distance measured as d with error sigma lies in [lo, hi).
double InBand(double d, double sigma, double lo, double hi) {
    const double s = std::sqrt(2.0) * sigma;
    return 0.5 * (std::erfc((lo - d) / s) - std::erfc((hi - d) / s));
}

// The coordinate error assumed for an atom: its rms displacement sqrt(B / 8 pi^2) times d_min / 10 A, a scale
// that gives ~0.1 A at B 20 A^2 and 2 A, typical of refined structures.
double CoordinateError(double b_iso, double d_min) {
    return std::sqrt(std::max(b_iso, 0.0) / (8 * PI * PI)) * d_min / 10.0;
}

struct Ligand {
    gemmi::Position pos;
    double distance;
    double sigma;                // the distance's error from the two atoms' coordinate errors
    double valence[BVS_CATIONS]; // its bond valence to each cation, occupancy weighted
    AtomChemistry chem;
    bool nitrogen;
    bool n7_off_plane;           // a purine N7 the site is out of the plane of, or at the wrong distance from its O6/N6
};

// The score of species `s` at `site`, with the N and O neighbours `ligands`. A plain sum of terms in log units:
// the first shell's count against the species' coordination number and its distances against the species' own,
// contacts the species does not make, the shell's chemistry, the bond-valence sum against the charge with its
// vector sum, and - where the maps are there - the electrons the difference map says the label lacks and the
// anomalous signal. A neighbour's coordinate error widens the bands it is counted in and the distance it is
// compared with, so that where the coordinates are poor the species' shells overlap and the score cannot tell
// them apart, rather than telling them apart wrongly. `label`: the label's species, -1 for one not scored here.
// What the model's waters show (SolventSite::electrons): the median and the rms spread of the log; NaN with too few.
struct WaterElectrons {
    double median = NAN, spread = NAN;
};
constexpr int WATER_STATISTICS_MIN = 10;

// `expected_electrons`: per species, the electrons a site of it shows on the scale of SolventSite::electrons, an ion
// at full occupancy; `waters`: what the model's waters show - many partly occupied or disordered - for the water.
double ScoreSpecies(int s, const SolventSite &site, const std::vector<Ligand> &ligands,
                    const std::array<double, SOLVENT_SPECIES> &expected_electrons, const WaterElectrons &waters) {
    const SpeciesRule &rule = RULES[s];
    const auto species = static_cast<SolventSpecies>(s);
    double n = 0, distance = 0, too_close = 0, excluded = 0;
    double donors_in = 0, acceptors_in = 0;   // donor-only and acceptor-only in the first shell
    double hbond_donors = 0, hbond_acceptors = 0;   // the same at 2.6-3.4 A, a water's hydrogen bonds
    double n7_off = 0;   // first-shell N7 a divalent ion could not hold
    for (const Ligand &l : ligands) {
        const double e = l.sigma;
        const double p = InBand(l.distance, BAND_EDGE, rule.lo - e, rule.hi + e);
        n += p;
        if (rule.esd > 0) {
            const double ideal = l.nitrogen ? rule.ideal_n : rule.ideal;
            const double esd2 = std::pow(l.nitrogen ? rule.esd_n : rule.esd, 2) + e * e;
            distance -= p * (l.distance - ideal) * (l.distance - ideal) / (2 * esd2);
        }
        too_close += InBand(l.distance, BAND_EDGE, 0.0, rule.too_close - e);
        if (rule.exclude_hi > 0)
            excluded += InBand(l.distance, BAND_EDGE, rule.exclude_lo + e, rule.exclude_hi);
        const bool donor_only = l.chem.donor && !l.chem.acceptor;
        const bool acceptor_only = l.chem.acceptor && !l.chem.donor;
        donors_in += donor_only ? p : 0;
        acceptors_in += acceptor_only ? p : 0;
        const double hb = InBand(l.distance, BAND_EDGE, 2.6 - e, SHELL_RADIUS + e);
        hbond_donors += donor_only ? hb : 0;
        hbond_acceptors += acceptor_only ? hb : 0;
        n7_off += l.n7_off_plane ? p : 0;
    }

    double score = 0;
    // the first shell
    if (species == SolventSpecies::Water) {
        if (n > rule.cn)
            score -= (n - rule.cn) * (n - rule.cn) / (2 * rule.cn_sigma * rule.cn_sigma);
        if (n < 1)
            score -= BAD_CONTACT * (1 - n);
    } else if (n < rule.cn && (species == SolventSpecies::Mg || species == SolventSpecies::Na ||
                               species == SolventSpecies::Ca || species == SolventSpecies::Cl))
        score -= MISSING_LIGAND * (rule.cn - n);
    else
        score -= (n - rule.cn) * (n - rule.cn) / (2 * rule.cn_sigma * rule.cn_sigma);
    score += distance;
    score -= BAD_CONTACT * (too_close + excluded);
    if (rule.charge == 2)
        score -= BAD_CONTACT * n7_off;

    // the chemistry
    if (rule.charge > 0) {
        score -= BAD_CONTACT * donors_in;              // an N-H does not point at a cation
        score -= BAD_CONTACT * site.positive_contacts;
    } else if (rule.charge < 0) {
        score -= BAD_CONTACT * acceptors_in;           // a halide is surrounded by donors
        score -= BAD_CONTACT * site.negative_contacts;
    } else {
        // a water donates two hydrogen bonds and accepts two
        score -= BAD_CONTACT * (std::max(0.0, hbond_donors - 2) + std::max(0.0, hbond_acceptors - 2));
    }

    // the bond valence, its expected spread widened by the ligands' distance errors
    if (rule.bvs >= 0) {
        double var = 0.25 * rule.charge * 0.25 * rule.charge;
        for (const Ligand &l : ligands)
            var += std::pow(l.valence[rule.bvs] * l.sigma / BV_B, 2);
        const double v = site.bvs[rule.bvs] - rule.charge;
        score -= v * v / (2 * var);
        // nVECSUM an outlier above 0.23, following CheckMyMetal, Zheng et al. (2017) Acta Cryst. D73, 223-233
        score -= 0.5 * std::pow(site.nvecsum[rule.bvs] / 0.23, 2);
    }

    // the electrons: the density at the site against what the species gives at the environment's B, an ion at full
    // occupancy, a water as the model's waters show - many partly occupied or disordered; log-normal in the waters'
    // spread (an ion's occupancy and order vary as much), at least the neighbours' calibration's. An ion as dense as a
    // full 10-electron scatterer is so more than a typical water, as Echols et al. (2014) flag a light ion by a
    // 2mFo-DFc above the waters' mean; where the density is weak the electrons tell little, and the shell decides.
    if (std::isfinite(site.electrons) && expected_electrons[s] > 0) {
        const bool known = std::isfinite(waters.median);
        const double expected = known && species == SolventSpecies::Water ? waters.median : expected_electrons[s];
        const double spread = known ? std::max(site.electron_spread, waters.spread) : site.electron_spread;
        const double q = std::max(site.electrons / expected, 0.05);
        score -= std::log(q) * std::log(q) / (2 * spread * spread);
    }
    // the anomalous signal against the species' f'' at the wavelength, Gaussian in the map's rms
    if (std::isfinite(site.anomalous_sigma) && std::isfinite(site.anomalous_expected[s]))
        score -= 0.5 * std::pow(site.anomalous_sigma - site.anomalous_expected[s], 2);
    return score;
}

std::string AtomName(const gemmi::Residue &r, const gemmi::Atom &a, const std::string &chain) {
    return a.name + " " + r.name + " " + chain + r.seqid.str();
}

}  // namespace

const char *SolventSpeciesName(SolventSpecies s) {
    switch (s) {
        case SolventSpecies::Water: return "HOH";
        case SolventSpecies::Mg: return "MG";
        case SolventSpecies::Na: return "NA";
        case SolventSpecies::K: return "K";
        case SolventSpecies::Ca: return "CA";
        case SolventSpecies::Cl: return "CL";
        case SolventSpecies::NH4: break;
    }
    return "NH4";
}

AtomChemistry ChemistryOf(const gemmi::MonLib &monlib, const gemmi::Residue &residue, const gemmi::Atom &atom,
                          bool n_terminus) {
    AtomChemistry c;
    if (!Polar(atom.element.elem))
        return c;
    if (IsWater(residue)) {
        c.donor = c.acceptor = true;
        return c;
    }
    const std::string &r = residue.name, &n = atom.name;
    char hb = 'B';   // an N or O the dictionary does not describe: either
    auto cc = monlib.monomers.find(r);
    if (cc != monlib.monomers.end()) {
        auto it = cc->second.find_atom(n);
        if (it != cc->second.atoms.end()) {
            auto e = monlib.ener_lib.atoms.find(it->chem_type);
            if (e != monlib.ener_lib.atoms.end())
                hb = e->second.hb_type;
        }
    }
    c.donor = hb == 'D' || hb == 'B';
    c.acceptor = hb == 'A' || hb == 'B';
    c.positive = (r == "LYS" && n == "NZ") || (r == "ARG" && (n == "NE" || n == "NH1" || n == "NH2")) ||
                 (n_terminus && n == "N");
    c.negative = (r == "ASP" && (n == "OD1" || n == "OD2")) || (r == "GLU" && (n == "OE1" || n == "OE2")) ||
                 n == "OXT" || (n == "O" && residue.find_atom("OXT", '*') != nullptr) ||
                 n == "OP1" || n == "OP2" || n == "OP3";
    c.flip_ambiguous = (r == "HIS" && (n == "ND1" || n == "NE2")) || (r == "ASN" && (n == "OD1" || n == "ND2")) ||
                       (r == "GLN" && (n == "OE1" || n == "NE2"));
    if (c.flip_ambiguous)
        c.donor = c.acceptor = true;
    return c;
}

bool SolventLabelIsWater(const SolventSite &s) {
    return LabelSpecies(s.label) == static_cast<int>(SolventSpecies::Water);
}

bool SolventLabelScored(const SolventSite &s) {
    const int label = LabelSpecies(s.label);
    return label >= 0 && label != static_cast<int>(SolventSpecies::NH4);
}

bool SolventLabelAgrees(const SolventSite &s) {
    return LabelSpecies(s.label) == static_cast<int>(s.best);
}

std::vector<SolventSite> AnalyseSolventSites(const gemmi::Structure &st, double d_min, const SolventMaps &maps,
                                             size_t nthreads) {
    std::vector<SolventSite> out;
    if (st.models.empty())
        return out;
    // per site, scored once the waters' electrons are known: its N and O neighbours and each species' electrons
    std::vector<std::vector<Ligand>> site_ligands;
    std::vector<std::array<double, SOLVENT_SPECIES>> site_expected;
    gemmi::Model model = st.models[0];   // NeighborSearch takes it non-const
    gemmi::UnitCell cell = st.cell;
    cell.set_cell_images_from_spacegroup(st.find_spacegroup());
    gemmi::NeighborSearch ns(model, cell, SEARCH_RADIUS);
    ns.populate(false);
    const gemmi::MonLib &monlib = EmbeddedMonomerLibrary();

    // the first residue of each polypeptide, whose N is charged
    std::vector<const gemmi::Residue *> n_termini;
    for (const gemmi::Chain &ch : model.chains) {
        const gemmi::ConstResidueSpan polymer = ch.get_polymer();
        if (!polymer.empty() && gemmi::find_tabulated_residue(polymer.front().name).is_amino_acid())
            n_termini.push_back(&polymer.front());
    }
    auto is_n_terminus = [&](const gemmi::Residue &r) {
        return std::find(n_termini.begin(), n_termini.end(), &r) != n_termini.end();
    };
    auto map_sigma = [](const gemmi::Grid<float> *g, double rms, const gemmi::Position &p) {
        return g != nullptr && rms > 0 ? g->interpolate_value(p) / rms : NAN;
    };
    const bool density = maps.map != nullptr && d_min > 0;
    RadialDensity radial(d_min);
    if (density) {
        std::vector<std::pair<Scatterer, double>> wanted;
        for (const gemmi::Chain &ch : model.chains)
            for (const gemmi::Residue &r : ch.residues)
                for (const gemmi::Atom &a : r.atoms)
                    wanted.emplace_back(ScattererOf(r, a), a.b_iso);
        radial.Prepare(wanted, nthreads);
    }
    // The model's atoms near the one at (ci, ri, ai) at `own` - in `near`, AtomsNear's list - other than it and its
    // coinciding copies, as they scatter.
    auto others_of = [&](size_t ci, size_t ri, size_t ai, const gemmi::Position &own,
                         const std::vector<ModelNeighbour> &near) {
        std::vector<Scattering> out;
        for (const ModelNeighbour &nb : near) {
            if (nb.chain == static_cast<int>(ci) && nb.residue == static_cast<int>(ri) && nb.atom == static_cast<int>(ai) &&
                nb.pos.dist(own) < 0.5)
                continue;
            const gemmi::Residue &r = model.chains[nb.chain].residues[nb.residue];
            const gemmi::Atom &a = r.atoms[nb.atom];
            out.push_back({nb.pos, a.occ, &radial.Table(ScattererOf(r, a), a.b_iso)});
        }
        return out;
    };
    auto tails = [&](const std::vector<Scattering> &others, const gemmi::Position &p) {
        double sum = 0;
        for (const Scattering &o : others)
            sum += o.occ * RadialDensity::At(*o.density, p.dist(o.pos));
        return sum;
    };
    // The map's scale, per electron (log), from the model's macromolecule: each C, N and O's height in 2mFo-DFc over
    // what the model puts there, the atom itself at its B and occupancy and the tails of its neighbours' - its
    // bonded neighbours and its riding hydrogens among them. Not from the waters: their occupancy is the model's
    // guess, and the depositions' waters show 0.3-0.5 of the height their model gives, the macromolecule's atoms
    // 0.8-1.
    std::map<const gemmi::Atom *, double> atom_calibration;
    std::pair<double, double> global_calibration{NAN, NAN};
    if (density) {
        std::vector<std::array<size_t, 3>> atoms;
        for (size_t ci = 0; ci < model.chains.size(); ci++)
            for (size_t ri = 0; ri < model.chains[ci].residues.size(); ri++) {
                const gemmi::Residue &r = model.chains[ci].residues[ri];
                if (IsWater(r) || IonAtom(r) != nullptr)
                    continue;
                for (size_t ai = 0; ai < r.atoms.size(); ai++) {
                    const gemmi::Atom &a = r.atoms[ai];
                    if ((a.element == gemmi::El::C || Polar(a.element.elem)) && a.occ == 1 && !a.altloc)
                        atoms.push_back({ci, ri, ai});
                }
            }
        std::vector<double> v(atoms.size());
        ParallelFor(static_cast<int>(atoms.size()), nthreads, [&](int k) {
            const auto [ci, ri, ai] = atoms[k];
            const gemmi::Residue &r = model.chains[ci].residues[ri];
            const gemmi::Atom &a = r.atoms[ai];
            const double model_height =
                    radial(ScattererOf(r, a), a.b_iso, 0.0) +
                    tails(others_of(ci, ri, ai, a.pos, AtomsNear(ns, model, a.pos, TAIL_RADIUS)), a.pos);
            v[k] = maps.map->interpolate_value(a.pos) / model_height;
        });
        std::vector<double> c;
        for (size_t k = 0; k < atoms.size(); k++)
            if (v[k] > 0) {
                const auto [ci, ri, ai] = atoms[k];
                atom_calibration[&model.chains[ci].residues[ri].atoms[ai]] = std::log(v[k]);
                c.push_back(std::log(v[k]));
            }
        global_calibration = MedianSpread(c);
    }
    // f'' by Cromer & Liberman (1981) Acta Cryst. A37, 267-268, corrected as Kissel & Pratt (1990) Acta Cryst. A46,
    // 170-175 (GEMMI's fprime).
    // The anomalous map's height per unit of f'' at B (a point scatterer's), least squares through the origin on the
    // model's S and P - the scatterers there are that sit where the model puts them; NaN with fewer than three.
    double fpp[SOLVENT_SPECIES];
    double anomalous_scale = NAN;
    if (maps.anomalous != nullptr && maps.anomalous_rms > 0 && maps.wavelength_A > 0 && d_min > 0) {
        const double energy = 12398.4198 / maps.wavelength_A;   // eV
        for (int k = 0; k < SOLVENT_SPECIES; k++)
            gemmi::cromer_liberman(static_cast<int>(SCATTERERS[k].el), energy, &fpp[k]);
        double fpp_s = 0, fpp_p = 0;
        gemmi::cromer_liberman(16, energy, &fpp_s);
        gemmi::cromer_liberman(15, energy, &fpp_p);
        double num = 0, den = 0;
        int n = 0;
        for (const gemmi::Chain &ch : model.chains)
            for (const gemmi::Residue &r : ch.residues)
                for (const gemmi::Atom &a : r.atoms)
                    if (a.element == gemmi::El::S || a.element == gemmi::El::P) {
                        const double unit = a.occ * (a.element == gemmi::El::S ? fpp_s : fpp_p) *
                                            PeakHeight([](double) { return 1.0; }, a.b_iso, d_min);
                        num += maps.anomalous->interpolate_value(a.pos) * unit;
                        den += unit * unit;
                        n++;
                    }
        if (n >= 3 && den > 0 && num > 0)
            anomalous_scale = num / den;
    }

    // the sites, in model order, each measured on a thread of its own
    std::vector<std::array<size_t, 3>> where;
    for (size_t ci = 0; ci < model.chains.size(); ci++)
        for (size_t ri = 0; ri < model.chains[ci].residues.size(); ri++) {
            const gemmi::Residue &res = model.chains[ci].residues[ri];
            if (IsWater(res) || IonAtom(res) != nullptr)
                for (size_t ai = 0; ai < res.atoms.size(); ai++)
                    if (!res.atoms[ai].is_hydrogen())
                        where.push_back({ci, ri, ai});
        }
    out.resize(where.size());
    site_ligands.resize(where.size());
    site_expected.resize(where.size());
    ParallelFor(static_cast<int>(where.size()), nthreads, [&](int index) {
        const auto [ci, ri, ai] = where[index];
        const gemmi::Residue &res = model.chains[ci].residues[ri];
        const gemmi::Atom &a = res.atoms[ai];
        SolventSite &site = out[index];
        site.label = res.name;
        site.chain = model.chains[ci].name;
        site.seqid = res.seqid.str();
        site.altloc = a.altloc;
        site.pos = a.pos;
        site.occupancy = a.occ;
        site.b_iso = a.b_iso;
        site.sigma_xyz = CoordinateError(a.b_iso, d_min);
        site.min_ligand_distance = NAN;
        site.nearest_positive = site.nearest_negative = site.nearest_ion = INFINITY;
        site.map_sigma = map_sigma(maps.map, maps.map_rms, a.pos);
        site.difference_sigma = map_sigma(maps.difference, maps.difference_rms, a.pos);
        site.anomalous_sigma = map_sigma(maps.anomalous, maps.anomalous_rms, a.pos);
        const std::vector<ModelNeighbour> near = AtomsNear(ns, model, a.pos, SEARCH_RADIUS + SOLVENT_PEAK_REACH);
        // The environment's B, and the map's scale from the macromolecule's atoms nearby where there are
        // enough of them.
        double b_sum = 0;
        int b_n = 0;
        std::vector<double> calibration;
        for (const ModelNeighbour &nb : near) {
            const gemmi::Atom &pa = model.chains[nb.chain].residues[nb.residue].atoms[nb.atom];
            if (nb.distance > SEARCH_RADIUS || (nb.chain == static_cast<int>(ci) && nb.residue == static_cast<int>(ri)))
                continue;
            if (const auto it = atom_calibration.find(&pa); it != atom_calibration.end())
                calibration.push_back(it->second);
            if (Polar(pa.element.elem) && nb.distance < BVS_RADIUS && (!a.altloc || !pa.altloc || a.altloc == pa.altloc)) {
                b_sum += pa.b_iso;
                b_n++;
            }
        }
        site.b_env = b_n > 0 ? b_sum / b_n : NAN;
        const double b_env = b_n > 0 ? site.b_env : a.b_iso;
        const double copies = a.occ * SiteSymmetry(cell, a.pos).multiplicity;
        const auto [log_k, spread] = static_cast<int>(calibration.size()) >= CALIBRATION_NEIGHBOURS
                                     ? MedianSpread(calibration) : global_calibration;
        // The site's own density fitted with its neighbours' held, at the environment's B; an ion, whose
        // label's restraints may hold it off its density, where that fit puts it.
        const bool ion = !IsWater(res);
        gemmi::Position q = a.pos;
        site.electrons = site.electron_spread = NAN;
        if (density && std::isfinite(log_k) && copies > 0) {
            const LocalFit fit = FitSite(*maps.map, std::exp(log_k), a.pos, b_env, others_of(ci, ri, ai, a.pos, near),
                                         radial, ion);
            q = fit.p;
            site.electrons = 10.0 * fit.a / copies;
            site.electron_spread = std::max(ELECTRON_SPREAD_MIN, spread);
        }
        site.measured_at = q;
        site.peak_shift = q.dist(a.pos);

        std::vector<Ligand> ligands;
        std::vector<gemmi::Position> bound_cations;   // other cations closer than 2.6 A
        std::vector<gemmi::Position> inner;   // the N and O closer than 2.6 A
        gemmi::Vec3 vec[BVS_CATIONS];
        for (ModelNeighbour nb : near) {
            // the site itself, its own hydrogens and its other conformers; its symmetry mates stay
            if (nb.chain == static_cast<int>(ci) && nb.residue == static_cast<int>(ri) &&
                (nb.atom != static_cast<int>(ai) || nb.pos.dist(a.pos) < 0.5))
                continue;
            nb.distance = q.dist(nb.pos);
            if (!(nb.distance < SEARCH_RADIUS))
                continue;
            const gemmi::Chain &pch = model.chains[nb.chain];
            const gemmi::Residue &pr = pch.residues[nb.residue];
            const gemmi::Atom &pa = pr.atoms[nb.atom];
            if (pa.is_hydrogen() || (a.altloc && pa.altloc && a.altloc != pa.altloc))
                continue;
            const double d = nb.distance;
            const int charge = IonCharge(pr);
            if (charge != 0) {
                site.ions_within_6++;
                site.ions_within_42 += d < 4.2;
                site.nearest_ion = std::min(site.nearest_ion, d);
                if (charge > 0) {
                    site.nearest_positive = std::min(site.nearest_positive, d);
                    site.positive_contacts += d < LIKE_CHARGE_DISTANCE;
                    if (d < 2.6)
                        bound_cations.push_back(nb.pos);
                } else {
                    site.nearest_negative = std::min(site.nearest_negative, d);
                    site.negative_contacts += d < HALIDE_NEGATIVE_DISTANCE;
                }
                continue;
            }
            const AtomChemistry chem = ChemistryOf(monlib, pr, pa, is_n_terminus(pr));
            if (chem.positive) {
                site.nearest_positive = std::min(site.nearest_positive, d);
                site.positive_contacts += d < LIKE_CHARGE_DISTANCE;
            }
            if (chem.negative) {
                site.nearest_negative = std::min(site.nearest_negative, d);
                site.negative_contacts += d < HALIDE_NEGATIVE_DISTANCE;
            }
            if (!Polar(pa.element.elem))
                continue;

            if (d < 1.9) site.n_below_19++;
            else if (d < 2.3) site.n_19_23++;
            else if (d < 2.6) site.n_23_26++;
            else if (d < 3.2) site.n_26_32++;
            else if (d < 3.4) site.n_32_34++;
            else if (d < 3.6) site.n_34_36++;
            if (d < SHELL_RADIUS) {
                if (chem.donor && chem.acceptor)
                    site.donor_acceptors++;
                else if (chem.donor)
                    site.donors++;
                else if (chem.acceptor)
                    site.acceptors++;
                site.neighbours += fmt::format("{}{} {:.2f}", site.neighbours.empty() ? "" : ", ",
                                               AtomName(pr, pa, pch.name), d);
            }
            if (d < 2.6)
                inner.push_back(nb.pos);

            Ligand l{};
            l.pos = nb.pos;
            l.distance = d;
            const double e = CoordinateError(pa.b_iso, d_min);
            l.sigma = std::sqrt(site.sigma_xyz * site.sigma_xyz + e * e);
            l.chem = chem;
            l.nitrogen = pa.element == gemmi::El::N;
            l.n7_off_plane = false;
            if (pa.name == "N7" && d < 2.6) {
                // the purine's ring, as the image the neighbour is, and its O6 or N6
                std::vector<gemmi::Atom> ring;
                const gemmi::Atom *o6n6 = nullptr;
                for (const gemmi::Atom &ra : pr.atoms) {
                    static const std::set<std::string> RING = {"N9", "C8", "N7", "C5", "C6", "N1", "C2", "N3", "C4"};
                    if (RING.count(ra.name) && (!a.altloc || !ra.altloc || ra.altloc == a.altloc)) {
                        ring.push_back(ra);
                        ring.back().pos = gemmi::Position(nb.image.apply(ra.pos));
                    }
                    if (ra.name == "O6" || ra.name == "N6")
                        o6n6 = &ra;
                }
                if (ring.size() == 9 && o6n6 != nullptr) {
                    site.n7_ligands++;
                    std::vector<gemmi::Atom *> ptrs;
                    for (gemmi::Atom &ra : ring)
                        ptrs.push_back(&ra);
                    const double plane = std::fabs(gemmi::get_distance_from_plane(q, gemmi::find_best_plane(ptrs)));
                    const double side = q.dist(gemmi::Position(nb.image.apply(o6n6->pos)));
                    l.n7_off_plane = plane > N7_PLANE_MAX + l.sigma || side < N7_O6N6_MIN - l.sigma ||
                                     side > N7_O6N6_MAX + l.sigma;
                    site.n7_off_plane += l.n7_off_plane;
                }
            }
            if (d < BVS_RADIUS)
                for (int c = 0; c < BVS_CATIONS; c++) {
                    // Following Brown & Altermatt (1985) Acta Cryst. B41, 244-247, occupancy weighted as
                    // Muller et al. (2003) Acta Cryst. D59, 32-37
                    l.valence[c] = pa.occ * std::exp(((l.nitrogen ? R0_N[c] : R0_O[c]) - d) / BV_B);
                    site.bvs[c] += l.valence[c];
                    vec[c] += (nb.pos - q) * (l.valence[c] / d);
                }
            ligands.push_back(l);
        }
        // A site bound to a cation shares that cation's other ligands as neighbours: the cis waters of a
        // Mg(H2O)6 are 2.9 A from each other, and are no hydrogen bonds of theirs. The scores leave them out.
        std::erase_if(ligands, [&](const Ligand &l) {
            for (const gemmi::Position &c : bound_cations)
                if (l.pos.dist(c) < 2.6)
                    return true;
            return false;
                });
                for (int c = 0; c < BVS_CATIONS; c++)
                    site.nvecsum[c] = site.bvs[c] > 0 ? vec[c].length() / site.bvs[c] : 0;
                for (size_t i = 0; i < inner.size(); i++)
                    for (size_t j = i + 1; j < inner.size(); j++)
                        site.min_ligand_distance = std::isnan(site.min_ligand_distance)
                                                   ? inner[i].dist(inner[j])
                                                   : std::min(site.min_ligand_distance, inner[i].dist(inner[j]));

                // The electrons each species would show on the scale of the fit (its density's height over a water's at
                // the environment's B), and the anomalous signal each would give. A site on a special position counts
                // its coinciding copies.
                std::array<double, SOLVENT_SPECIES> expected_electrons{};
                if (std::isfinite(site.electrons))
                    for (int k = 0; k < SOLVENT_SPECIES; k++)
                        expected_electrons[k] = 10.0 * radial(SCATTERERS[k], b_env, 0.0) / radial(SCATTERERS[0], b_env, 0.0);
                for (int k = 0; k < SOLVENT_SPECIES; k++)
                    site.anomalous_expected[k] =
                            std::isfinite(anomalous_scale)
                            ? anomalous_scale * fpp[k] * copies * PeakHeight([](double) { return 1.0; }, b_env, d_min) /
                              maps.anomalous_rms
                            : NAN;
                if (std::isfinite(anomalous_scale))
                    site.anomalous_sigma = map_sigma(maps.anomalous, maps.anomalous_rms, q);

                site_ligands[index] = std::move(ligands);
                site_expected[index] = expected_electrons;
    });

    WaterElectrons waters;
    std::vector<double> w;
    for (const SolventSite &site : out)
        if (SolventLabelIsWater(site) && site.electrons > 0)
            w.push_back(std::log(site.electrons));
    if (static_cast<int>(w.size()) >= WATER_STATISTICS_MIN) {
        const auto [m, spread] = MedianSpread(w);
        waters = {std::exp(m), spread};
    }
    for (size_t i = 0; i < out.size(); i++) {
        SolventSite &site = out[i];
        for (int s = 0; s < SOLVENT_SPECIES; s++)
            site.score[s] = ScoreSpecies(s, site, site_ligands[i], site_expected[i], waters);
        // NH4+ is reported, not decided: water, Mg, Na, K, Ca, Cl
        int best = 0, second = -1;
        for (int s = 1; s < static_cast<int>(SolventSpecies::NH4); s++)
            if (site.score[s] > site.score[best]) {
                second = best;
                best = s;
            } else if (second < 0 || site.score[s] > site.score[second])
                second = s;
        site.best = static_cast<SolventSpecies>(best);
        site.ambiguous = site.score[best] - site.score[second] < SOLVENT_AMBIGUOUS_MARGIN;
    }
    return out;
}