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

#include "SolventBuild.h"

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

#include <spdlog/fmt/fmt.h>

#include "gemmi/asumask.hpp"
#include "gemmi/math.hpp"
#include "gemmi/neighbor.hpp"

#include "MonomerLibrary.h"
#include "RefinementRestraints.h"   // AtomsNear

namespace {

// The ordered-solvent defaults of phenix.refine: peaks of mFo-DFc from 3 sigma, hydrogen-bond partners to 3.2 A
// (taken here to 3.4 A, the outer edge of the coordinate error).
// Following Afonine et al. (2012) Acta Cryst. D68, 352-367
constexpr double PEAK_SIGMA = 3.0;
constexpr double PARTNER_MAX = 3.4;            // A
// A water no nearer than 2.2 A to an N or O (one that close is a cation's ligand or an error), nor than 3.0 A
// to a C or any other atom that makes no hydrogen bond; Coot's Find Waters keeps waters 2.4-3.2 A from the
// model, Emsley et al. (2010) Acta Cryst. D66, 486-501, here a little wider for the peak's own position error.
constexpr double POLAR_CLASH = 2.2;            // A
constexpr double APOLAR_CLASH = 3.0;           // A
// A: a water this near a metal is its ligand or in its shell (K+ holds its waters at up to 3.2 A, Cat_Wiz's band)
constexpr double METAL_LIGAND_MAX = 3.2;
constexpr double HALIDE_PARTNER_MIN = 2.9;     // A: a water donates to a halide from here
// Nearer than ON_ATOM to a model atom, the peak is that atom's own residual (its B, its anisotropy) and no candidate;
// out to BONDED it is bonded to it (a missing or displaced atom), not a water.
constexpr double ON_ATOM = 1.2, BONDED = 1.9;  // A
// A difference peak beside a model atom that sits in negative difference density: the atom is displaced.
// A water in negative density at its centre has more electrons modelled than the data show.
// Following Echols et al. (2014) Acta Cryst. D70, 1104-1114 (their -3 sigma)
constexpr double NEGATIVE_SIGMA = -3.0;
constexpr double NEGATIVE_REACH = 2.6;         // A: the model atom the peak is beside
// Peaks of one molecule: local maxima this far apart are bonded atoms, nearer they are one peak.
constexpr double CLUSTER_MIN = 1.1, CLUSTER_MAX = 1.9;   // A
// A tetrahedral cluster: three or more peaks 1.3-1.7 A from one (sulfate, phosphate).
constexpr double TETRAHEDRAL_MIN = 1.3, TETRAHEDRAL_MAX = 1.7;   // A
// The density a peak is measured in for its shape: the points above the cutoff within this distance.
constexpr double SHAPE_RADIUS = 1.6;           // A
// rms along the longest axis over the shortest; the peaks at deposited waters of five refine-arm sets (stripped and
// rebuilt) reach 3.0 at their 90th percentile, the measure being coarse on a few grid points.
constexpr double ELONGATION_MAX = 3.0;
// A peak this near the model's truncated residues' last atoms is where their missing atoms would be.
constexpr double TRUNCATED_REACH = 3.0;        // A

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

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

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

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

// Each grid point within `radius` of `p`, its position and value, the cell's periodicity taken care of.
template <class F>
void ForPointsNear(const gemmi::Grid<float> &g, const gemmi::Position &p, double radius, F &&f) {
    const gemmi::Fractional fp = g.unit_cell.fractionalize(p);
    const int u0 = static_cast<int>(std::round(fp.x * g.nu)), v0 = static_cast<int>(std::round(fp.y * g.nv)),
              w0 = static_cast<int>(std::round(fp.z * g.nw));
    const int du = static_cast<int>(std::ceil(radius / g.spacing[0])) + 1;
    const int dv = static_cast<int>(std::ceil(radius / g.spacing[1])) + 1;
    const int dw = static_cast<int>(std::ceil(radius / g.spacing[2])) + 1;
    for (int w = w0 - dw; w <= w0 + dw; w++)
        for (int v = v0 - dv; v <= v0 + dv; v++)
            for (int u = u0 - du; u <= u0 + du; u++) {
                const gemmi::Position q = g.get_position(u, v, w);
                if (q.dist(p) <= radius)
                    f(q, g.get_value(u, v, w), u, v, w);
            }
}

bool IsLocalMaximum(const gemmi::Grid<float> &g, int u, int v, int w) {
    const float value = g.get_value(u, v, w);
    for (int dw = -1; dw <= 1; dw++)
        for (int dv = -1; dv <= 1; dv++)
            for (int du = -1; du <= 1; du++)
                if ((du || dv || dw) && g.get_value(u + du, v + dv, w + dw) > value)
                    return false;
    return true;
}

struct Peak {
    gemmi::Position pos;
    double sigma;
};

// The local maxima of `g` (in rms `rms`) from `cutoff`, each symmetry-equivalent set once, strongest first.
std::vector<Peak> DifferencePeaks(const gemmi::Grid<float> &g, double rms, double cutoff) {
    std::vector<Peak> out;
    const std::vector<std::int8_t> asu = gemmi::get_asu_mask(g);
    for (int w = 0; w < g.nw; w++)
        for (int v = 0; v < g.nv; v++)
            for (int u = 0; u < g.nu; u++) {
                const size_t idx = g.index_q(u, v, w);
                if (asu[idx] != 0 || g.data[idx] < cutoff * rms || !IsLocalMaximum(g, u, v, w))
                    continue;
                out.push_back({g.get_position(u, v, w), g.data[idx] / rms});
            }
    std::sort(out.begin(), out.end(), [](const Peak &a, const Peak &b) { return a.sigma > b.sigma; });
    return out;
}

// The copy of `p` nearest to the model's own atoms (not their symmetry mates); false where no atom is within
// `reach`.
bool NearestCopy(gemmi::NeighborSearch &ns, const gemmi::UnitCell &cell, double reach, gemmi::Position &p) {
    gemmi::NeighborSearch::Mark *m = ns.find_nearest_atom(p, reach);
    if (m == nullptr)
        return false;
    const gemmi::FTransform t = ns.get_image_transformation(m->image_idx);
    const gemmi::Fractional fc = cell.fractionalize(p);
    const gemmi::Fractional fi = t.apply(cell.fractionalize(m->pos));
    const gemmi::Fractional shifted(fc.x - std::round(fc.x - fi.x), fc.y - std::round(fc.y - fi.y),
                                    fc.z - std::round(fc.z - fi.z));
    p = cell.orthogonalize(gemmi::FTransform(t.inverse()).apply(shifted));
    return true;
}

// The longest over the shortest rms axis of the density above `cutoff` (in map units) within SHAPE_RADIUS of
// `p`, each point weighted by how far it is above; 1 where there are too few points to tell.
double Elongation(const gemmi::Grid<float> &g, const gemmi::Position &p, double cutoff) {
    std::vector<std::pair<gemmi::Position, double>> points;
    double sum = 0;
    gemmi::Vec3 mean;
    ForPointsNear(g, p, SHAPE_RADIUS, [&](const gemmi::Position &q, float value, int, int, int) {
        if (value > cutoff) {
            points.emplace_back(q, value - cutoff);
            sum += value - cutoff;
            mean += q * (value - cutoff);
        }
    });
    if (points.size() < 6)
        return 1.0;
    mean /= sum;
    double c[6] = {0, 0, 0, 0, 0, 0};
    for (const auto &[q, wt] : points) {
        const gemmi::Vec3 d = gemmi::Vec3(q) - mean;
        c[0] += wt * d.x * d.x;
        c[1] += wt * d.y * d.y;
        c[2] += wt * d.z * d.z;
        c[3] += wt * d.x * d.y;
        c[4] += wt * d.x * d.z;
        c[5] += wt * d.y * d.z;
    }
    const gemmi::SMat33<double> m{c[0] / sum, c[1] / sum, c[2] / sum, c[3] / sum, c[4] / sum, c[5] / sum};
    std::array<double, 3> ev = m.calculate_eigenvalues();
    std::sort(ev.begin(), ev.end());
    return ev[0] > 0 ? std::sqrt(ev[2] / ev[0]) : INFINITY;
}

// The atoms of residues the dictionary describes with heavy atoms the model lacks, that the dictionary bonds to
// one of the missing atoms: the ends a truncated side chain stops at.
std::set<const gemmi::Atom *> TruncatedEnds(const gemmi::Model &model) {
    std::set<const gemmi::Atom *> out;
    const gemmi::MonLib &monlib = EmbeddedMonomerLibrary();
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues) {
            auto cc = monlib.monomers.find(r.name);
            if (cc == monlib.monomers.end() || r.is_water())
                continue;
            auto missing = [&](const std::string &name) {
                auto a = cc->second.find_atom(name);
                return a != cc->second.atoms.end() && !a->is_hydrogen() && name != "OXT" && name != "OP3" &&
                       r.find_atom(name, '*') == nullptr;
            };
            for (const gemmi::Restraints::Bond &b : cc->second.rt.bonds)
                for (int k = 0; k < 2; k++) {
                    const std::string &here = k == 0 ? b.id1.atom : b.id2.atom, &there = k == 0 ? b.id2.atom : b.id1.atom;
                    if (missing(there))
                        if (const gemmi::Atom *a = r.find_atom(here, '*'))
                            out.insert(a);
                }
        }
    return out;
}

}  // namespace

const char *SolventEditKindName(SolventEdit::Kind kind) {
    switch (kind) {
        case SolventEdit::Kind::Added: return "ADDED";
        case SolventEdit::Kind::Removed: return "REMOVED";
        case SolventEdit::Kind::NotSolvent: return "NOT_SOLVENT";
        case SolventEdit::Kind::Relabelled: return "RELABELLED";
    }
    return "?";
}

SolventRound PickSolvent(const gemmi::Structure &st, const SolventMaps &maps, const std::vector<SolventSite> &sites,
                         const std::set<std::string> &built, const std::set<std::string> &relabelled, double b_ceiling,
                         int round) {
    SolventRound out;
    if (st.models.empty() || maps.difference == nullptr || maps.map == nullptr || !(maps.difference_rms > 0) ||
        !(maps.map_rms > 0))
        return out;
    const gemmi::Grid<float> &diff = *maps.difference, &map = *maps.map;
    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, 5.0);
    ns.populate(false);
    const gemmi::MonLib &monlib = EmbeddedMonomerLibrary();
    const std::set<const gemmi::Atom *> truncated = TruncatedEnds(model);

    // The densest S and P of the model: a peak denser than every one of them is no light atom.
    double heaviest = -INFINITY;
    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)
                    heaviest = std::max<double>(heaviest, map.interpolate_value(a.pos) / maps.map_rms);
    if (!std::isfinite(heaviest))
        heaviest = INFINITY;

    // The waters the maps no longer support: any in negative difference density, and those built here also where they
    // have no partner left or their B reached the ceiling. phenix.refine's other filter, 2mFo-DFc below 1 sigma, is
    // left out: on the refine arm it raised R-free wherever it was tried - on the depositions' waters (five sets of
    // five) and on the waters built here (the next round rejected in five of six) - and so did taking out the
    // depositions' waters without a partner (1K4C's 38). A weak water still explains density.
    for (const SolventSite &s : sites) {
        if (!SolventLabelIsWater(s))
            continue;
        std::string reason;
        const std::string name = s.label + " " + s.chain + s.seqid;
        const bool own = built.count(name) > 0;
        if (s.difference_sigma <= NEGATIVE_SIGMA)
            reason = fmt::format("mFo-DFc {:.1f} sigma", s.difference_sigma);
        else if (own && s.n_below_19 + s.n_19_23 + s.n_23_26 + s.n_26_32 + s.n_32_34 == 0 && !(s.nearest_ion < PARTNER_MAX))
            reason = fmt::format("no N, O or ion within {} A", PARTNER_MAX);
        else if (own && s.b_iso >= b_ceiling - 0.5)
            reason = fmt::format("B {:.0f} at the refinement's ceiling", s.b_iso);
        if (reason.empty())
            continue;
        SolventEdit e;
        e.kind = SolventEdit::Kind::Removed;
        e.round = round;
        e.site = name;
        e.chain = s.chain;
        e.seqid = gemmi::SeqId(s.seqid);
        e.altloc = s.altloc;
        e.pos = s.pos;
        e.difference_sigma = s.difference_sigma;
        e.map_sigma = s.map_sigma;
        e.b_iso = s.b_iso;
        e.reason = reason;
        out.edits.push_back(e);
        out.removed++;
    }
    // The sites whose surroundings and density support another species than their label, not ambiguously: each
    // relabelled once (an alternative conformer, a site relabelled before, or one with no density of its own to tell
    // the species by, stays as it is).
    std::set<std::string> taken_out;
    for (const SolventEdit &e : out.edits)
        taken_out.insert(e.site);
    for (const SolventSite &s : sites) {
        const std::string name = s.label + " " + s.chain + s.seqid;
        if (!SolventLabelScored(s) || s.ambiguous || SolventLabelAgrees(s) || s.altloc || taken_out.count(name) ||
            relabelled.count(name) || !(s.electrons > 0))
            continue;
        const int best = static_cast<int>(s.best);
        int label = 0;
        while (label < SOLVENT_SPECIES && s.label != SolventSpeciesName(static_cast<SolventSpecies>(label)))
            label++;
        if (SolventLabelIsWater(s))
            label = static_cast<int>(SolventSpecies::Water);
        SolventEdit e;
        e.kind = SolventEdit::Kind::Relabelled;
        e.round = round;
        e.site = name;
        e.chain = s.chain;
        e.seqid = gemmi::SeqId(s.seqid);
        e.pos = s.pos;
        e.difference_sigma = s.difference_sigma;
        e.map_sigma = s.map_sigma;
        e.b_iso = s.b_iso;
        e.species = SolventSpeciesName(s.best);
        e.reason = fmt::format("{} {:.1f} against {} {:.1f}; shell {}/{}/{}, electrons {:.1f}, anomalous {:.1f} sigma",
                               e.species, s.score[best], s.label,
                               label < SOLVENT_SPECIES ? s.score[label] : NAN, s.n_19_23, s.n_23_26, s.n_26_32,
                               s.electrons, s.anomalous_sigma);
        out.edits.push_back(e);
        out.relabelled++;
    }

    // A removed water is no partner of a new one, nor in its way.
    std::vector<gemmi::Position> removed;
    for (const SolventEdit &e : out.edits)
        if (e.kind == SolventEdit::Kind::Removed)
            removed.push_back(e.pos);
    auto is_removed = [&](const gemmi::Position &p) {
        return std::any_of(removed.begin(), removed.end(), [&](const gemmi::Position &q) { return q.dist(p) < 0.01; });
    };

    // The difference map's peaks, strongest first.
    std::vector<gemmi::Position> added;
    for (Peak peak : DifferencePeaks(diff, maps.difference_rms, PEAK_SIGMA)) {
        if (!NearestCopy(ns, cell, PARTNER_MAX + 0.1, peak.pos))
            continue;   // nothing of the model near enough to hold it
        // On a special position (SPECIAL_POSITION_TOLERANCE), exactly on it, its copies sharing one water; nearer
        // than a hydrogen bond to its own copy otherwise, it would clash with itself.
        const SpecialPosition site = SiteSymmetry(cell, peak.pos);
        if (!site.ops.empty())
            peak.pos = site.OnSite(peak.pos);
        else if (cell.is_special_position(peak.pos, POLAR_CLASH) > 0) {
            out.refused["too close to its own symmetry copy"]++;
            continue;
        }
        SolventEdit e;
        e.round = round;
        e.pos = peak.pos;
        e.occupancy = 1.0 / site.multiplicity;
        e.difference_sigma = peak.sigma;
        e.map_sigma = map.interpolate_value(peak.pos) / maps.map_rms;
        e.elongation = Elongation(diff, peak.pos, PEAK_SIGMA * maps.difference_rms);
        auto refuse = [&](const std::string &why) { out.refused[why]++; };
        auto not_solvent = [&](const std::string &why) {
            e.kind = SolventEdit::Kind::NotSolvent;
            e.reason = why;
            out.edits.push_back(e);
            out.not_solvent++;
        };

        const std::vector<ModelNeighbour> near = AtomsNear(ns, model, peak.pos, PARTNER_MAX);
        const ModelNeighbour *nearest = nullptr;
        std::string clash, truncated_by;
        std::vector<std::string> partners;
        double partner_b = 0;
        for (const ModelNeighbour &nb : near) {
            const gemmi::Chain &ch = model.chains[nb.chain];
            const gemmi::Residue &r = ch.residues[nb.residue];
            const gemmi::Atom &a = r.atoms[nb.atom];
            if (a.is_hydrogen() || (r.is_water() && is_removed(a.pos)))
                continue;
            if (nearest == nullptr)
                nearest = &nb;
            const double d = nb.distance;
            const gemmi::El el = a.element.elem;
            if (d < TRUNCATED_REACH && truncated_by.empty() && truncated.count(&a))
                truncated_by = ResidueName(r, ch.name);
            bool partner = false;
            if (a.element.is_metal())
                partner = d < METAL_LIGAND_MAX;
            else if (Polar(el)) {
                if (d < POLAR_CLASH && clash.empty())
                    clash = AtomName(r, a, ch.name);
                const AtomChemistry chem = ChemistryOf(monlib, r, a, false);
                partner = d >= POLAR_CLASH && (chem.donor || chem.acceptor);
            } else if (IsHalogen(el))
                partner = d >= HALIDE_PARTNER_MIN;
            else if (d < APOLAR_CLASH && clash.empty())
                clash = AtomName(r, a, ch.name);
            if (partner) {
                partners.push_back(fmt::format("{} {:.2f}", AtomName(r, a, ch.name), d));
                partner_b += a.b_iso;
            }
        }
        if (nearest == nullptr)
            continue;
        const gemmi::Chain &nch = model.chains[nearest->chain];
        const gemmi::Residue &nres = nch.residues[nearest->residue];
        const gemmi::Atom &natom = nres.atoms[nearest->atom];

        if (nearest->distance < ON_ATOM) {
            refuse("on a model atom");
            continue;
        }
        // Not a solvent atom: reported, never filled.
        if (nearest->distance < BONDED) {
            not_solvent(fmt::format("{:.2f} A from {}: a missing or displaced atom of the model", nearest->distance,
                                    AtomName(nres, natom, nch.name)));
            continue;
        }
        if (nearest->distance < NEGATIVE_REACH &&
            diff.interpolate_value(nearest->pos) / maps.difference_rms <= NEGATIVE_SIGMA) {
            not_solvent(fmt::format("beside {} in negative density: the atom is displaced",
                                    AtomName(nres, natom, nch.name)));
            continue;
        }
        if (!truncated_by.empty()) {
            not_solvent(fmt::format("where the missing atoms of {} would be", truncated_by));
            continue;
        }
        int bonded = 0, tetrahedral = 0;
        ForPointsNear(diff, peak.pos, CLUSTER_MAX, [&](const gemmi::Position &q, float value, int u, int v, int w) {
            const double d = q.dist(peak.pos);
            if (d >= CLUSTER_MIN && value >= PEAK_SIGMA * maps.difference_rms && IsLocalMaximum(diff, u, v, w)) {
                bonded++;
                tetrahedral += d >= TETRAHEDRAL_MIN && d <= TETRAHEDRAL_MAX;
            }
        });
        if (tetrahedral >= 3) {
            not_solvent(fmt::format("the centre of {} peaks at {}-{} A: a tetrahedral group (sulfate, phosphate)",
                                    tetrahedral, TETRAHEDRAL_MIN, TETRAHEDRAL_MAX));
            continue;
        }
        if (bonded > 0) {
            not_solvent(fmt::format("{} other peak{} at {}-{} A: atoms of one molecule", bonded, bonded > 1 ? "s" : "",
                                    CLUSTER_MIN, CLUSTER_MAX));
            continue;
        }
        if (e.elongation > ELONGATION_MAX) {
            not_solvent(fmt::format("elongated ({:.1f}:1)", e.elongation));
            continue;
        }
        if (e.map_sigma > heaviest) {
            not_solvent(fmt::format("2mFo-DFc {:.1f} sigma, above every S and P of the model ({:.1f})", e.map_sigma,
                                    heaviest));
            continue;
        }

        // A water, where it has a partner and nothing in its way.
        if (!clash.empty()) {
            refuse("too close to the model");
            continue;
        }
        if (partners.empty()) {
            refuse("no hydrogen-bond partner");
            continue;
        }
        if (std::any_of(added.begin(), added.end(), [&](const gemmi::Position &q) {
                return cell.find_nearest_image(peak.pos, q, gemmi::Asu::Any).dist() < POLAR_CLASH;
            })) {
            refuse("too close to a water added beside it");
            continue;
        }
        e.kind = SolventEdit::Kind::Added;
        e.chain = nch.name;
        e.b_iso = partner_b / partners.size();
        e.reason = "partners";
        for (size_t k = 0; k < partners.size(); k++)
            e.reason += (k == 0 ? " " : ", ") + partners[k];
        out.edits.push_back(e);
        out.added++;
        added.push_back(peak.pos);
    }
    return out;
}

std::vector<InnerSphereDonor> InnerSphereDonors(const gemmi::Model &model, const gemmi::Residue &own, const gemmi::Atom &metal) {
    std::vector<InnerSphereDonor> donors;
    double widest = 0, sigma = 0;
    for (gemmi::El el : {gemmi::El::O, gemmi::El::N, gemmi::El::S}) {
        double d0 = 0;
        if (EmbeddedMetalDistance(metal.element.elem, el, 6, d0, sigma))
            widest = std::max(widest, d0 + INNER_SHELL_MARGIN + 0.2);
    }
    for (size_t c = 0; c < model.chains.size(); c++)
        for (size_t r = 0; r < model.chains[c].residues.size(); r++) {
            const gemmi::Residue &res = model.chains[c].residues[r];
            for (size_t k = 0; k < res.atoms.size(); k++) {
                const gemmi::Atom &d = res.atoms[k];
                if (&res != &own && d.occ > 0 &&
                    (d.element == gemmi::El::O || d.element == gemmi::El::N || d.element == gemmi::El::S))
                    if (const double dist = d.pos.dist(metal.pos); dist < widest)
                        donors.push_back({static_cast<int>(c), static_cast<int>(r), static_cast<int>(k), dist});
            }
        }
    auto ideal = [&](const InnerSphereDonor &d, int cn, double &value) {
        const gemmi::Atom &a = model.chains[d.chain].residues[d.residue].atoms[d.atom];
        return EmbeddedMetalDistance(metal.element.elem, a.element.elem, cn, value, sigma);
    };
    // judged at the coordination of those within the inner sphere of an octahedral site, as AddMetalLinks judges
    // a model's records at the coordination they claim
    int cn = 0;
    for (const InnerSphereDonor &d : donors) {
        double value = 0;
        cn += ideal(d, 6, value) && d.distance <= value + INNER_SHELL_MARGIN;
    }
    std::erase_if(donors, [&](const InnerSphereDonor &d) {
        double value = 0;
        return !ideal(d, std::max(cn, 1), value) || d.distance > value + INNER_SHELL_MARGIN;
    });
    return donors;
}

void AddMetalRecords(gemmi::Structure &st, const gemmi::Chain &chain, const gemmi::Residue &res, int &record) {
    const gemmi::Model &model = st.models.at(0);
    const gemmi::Atom &a = res.atoms.at(0);
    if (!a.element.is_metal())
        return;
    for (const InnerSphereDonor &d : InnerSphereDonors(model, res, a)) {
        const gemmi::Residue &dr = model.chains[d.chain].residues[d.residue];
        const gemmi::Atom &da = dr.atoms[d.atom];
        gemmi::Connection c;
        c.name = "rugnux_metalc" + std::to_string(++record);
        c.type = gemmi::Connection::MetalC;
        c.asu = gemmi::Asu::Same;
        c.partner1 = gemmi::AtomAddress(model.chains[d.chain].name, dr.seqid, dr.name, da.name, da.altloc);
        c.partner2 = gemmi::AtomAddress(chain.name, res.seqid, res.name, a.name);
        c.reported_distance = d.distance;
        st.connections.push_back(c);
    }
}

gemmi::El SolventSpeciesElement(const std::string &species) {
    return species == "HOH" ? gemmi::El::O : gemmi::Element(species).elem;
}

namespace {

// The site of `e` relabelled as e.species, its metal records replaced by its inner sphere's.
void Relabel(gemmi::Structure &st, const SolventEdit &e, int &record) {
    gemmi::Model &model = st.models.at(0);
    const std::string from = e.site.substr(0, e.site.find(' '));
    gemmi::Residue *res = nullptr;
    gemmi::Chain *chain = nullptr;
    for (gemmi::Chain &ch : model.chains)
        for (gemmi::Residue &r : ch.residues)
            if (ch.name == e.chain && r.seqid == e.seqid && r.name == from && r.atoms.size() == 1) {
                res = &r;
                chain = &ch;
            }
    if (res == nullptr)
        return;
    const gemmi::AtomAddress old_address(chain->name, res->seqid, res->name, res->atoms[0].name);
    std::erase_if(st.connections, [&](const gemmi::Connection &c) {
        return c.partner1 == old_address || c.partner2 == old_address;
    });
    // the subchain and entity of a residue of the species already in the model, where there is one
    const gemmi::Residue *like = nullptr;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            if (r.name == e.species && like == nullptr)
                like = &r;
    const bool water = e.species == "HOH";
    res->name = e.species;
    res->entity_type = water ? gemmi::EntityType::Water : gemmi::EntityType::NonPolymer;
    res->het_flag = 'H';
    res->subchain = like != nullptr ? like->subchain : "";
    res->entity_id = like != nullptr ? like->entity_id : "";
    gemmi::Atom &a = res->atoms[0];
    a.element = gemmi::Element(SolventSpeciesElement(e.species));
    a.name = water ? "O" : e.species;
    AddMetalRecords(st, *chain, *res, record);
}

}  // namespace

void ApplySolventEdits(gemmi::Structure &st, std::vector<SolventEdit> &edits) {
    gemmi::Model &model = st.models.at(0);
    int record = 0;
    for (const gemmi::Connection &c : st.connections)
        if (c.name.rfind("rugnux_metalc", 0) == 0)
            record = std::max(record, std::atoi(c.name.c_str() + 13));
    for (const SolventEdit &e : edits)
        if (e.kind == SolventEdit::Kind::Relabelled)
            Relabel(st, e, record);
    // A chain name can stand for several chains of the model (a polymer and its waters read as two); each is meant.
    for (const SolventEdit &e : edits) {
        if (e.kind != SolventEdit::Kind::Removed)
            continue;
        for (gemmi::Chain &ch : model.chains) {
            if (ch.name != e.chain)
                continue;
            for (gemmi::Residue &r : ch.residues)
                if (r.is_water() && r.seqid == e.seqid)
                    std::erase_if(r.atoms, [&](const gemmi::Atom &a) { return a.altloc == e.altloc && !a.is_hydrogen(); });
            std::erase_if(ch.residues, [](const gemmi::Residue &r) {
                return r.is_water() && std::none_of(r.atoms.begin(), r.atoms.end(), [](const gemmi::Atom &a) {
                    return !a.is_hydrogen();
                });
            });
        }
    }
    for (SolventEdit &e : edits) {
        if (e.kind != SolventEdit::Kind::Added)
            continue;
        // to the last chain of the name, numbered on from the last residue of any of them
        gemmi::Chain *ch = nullptr;
        int last = 0;
        const gemmi::Residue *water = nullptr;
        for (gemmi::Chain &c : model.chains) {
            if (c.name != e.chain)
                continue;
            ch = &c;
            for (const gemmi::Residue &r : c.residues) {
                if (r.seqid.num != gemmi::SeqId::OptionalNum::None)
                    last = std::max(last, static_cast<int>(r.seqid.num));
                if (r.is_water())
                    water = &r;
            }
        }
        if (ch == nullptr)
            continue;
        gemmi::Residue r;
        r.name = "HOH";
        r.seqid = gemmi::SeqId(last + 1, ' ');
        r.het_flag = 'H';
        r.entity_type = gemmi::EntityType::Water;
        if (water != nullptr) {
            r.subchain = water->subchain;
            r.entity_id = water->entity_id;
        }
        gemmi::Atom a;
        a.name = "O";
        a.element = gemmi::El::O;
        a.pos = e.pos;
        a.occ = static_cast<float>(e.occupancy);
        a.b_iso = static_cast<float>(e.b_iso);
        r.atoms.push_back(a);
        e.site = ResidueName(r, ch->name);
        e.seqid = r.seqid;
        ch->residues.push_back(r);
    }
}