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

#include <catch2/catch_all.hpp>

#include <array>
#include <cmath>
#include <string>
#include <vector>

#include <gemmi/dencalc.hpp>
#include <gemmi/it92.hpp>

#include "../image_analysis/structure_refinement/MonomerLibrary.h"
#include "../image_analysis/structure_refinement/RefinementRestraints.h"
#include "../image_analysis/structure_refinement/SolventSites.h"

namespace {

// A P1 box with nothing in it, large enough that no lattice translate comes near the middle.
gemmi::Structure empty_box(const char *group = "P 1", double edge = 40.0) {
    gemmi::Structure st;
    st.cell = gemmi::UnitCell(edge, edge, edge, 90, 90, 90);
    st.spacegroup_hm = group;
    st.setup_cell_images();
    st.models.emplace_back(1);
    st.models[0].chains.emplace_back("A");
    return st;
}

// One atom as a residue of its own, numbered after the last.
void add_atom(gemmi::Structure &st, const std::string &resname, const std::string &atom, gemmi::El el,
              const gemmi::Position &pos) {
    gemmi::Chain &ch = st.models[0].chains[0];
    gemmi::Residue r;
    r.name = resname;
    r.seqid = gemmi::SeqId(static_cast<int>(ch.residues.size()) + 1, ' ');
    gemmi::Atom a;
    a.name = atom;
    a.element = gemmi::Element(el);
    a.pos = pos;
    a.occ = 1.0f;
    a.b_iso = 20.0f;
    r.atoms.push_back(a);
    ch.residues.push_back(r);
}

const gemmi::Position CENTRE(20, 20, 20);

const SolventSite &site_labelled(const std::vector<SolventSite> &sites, const std::string &label) {
    for (const SolventSite &s : sites)
        if (s.label == label)
            return s;
    FAIL("no site labelled " << label);
    return sites.front();
}

}  // namespace

TEST_CASE("SolventSites_AtomsNear_symmetry_mate", "[Solvent]") {
    // an atom near the 2-fold axis of P2: its mate (-x, y, -z) is 1.70 A away
    gemmi::Structure st = empty_box("P 1 2 1", 30.0);
    const gemmi::Position p = st.cell.orthogonalize(gemmi::Fractional(0.02, 0.3, 0.02));
    add_atom(st, "HOH", "O", gemmi::El::O, p);
    gemmi::UnitCell cell = st.cell;
    cell.set_cell_images_from_spacegroup(st.find_spacegroup());
    gemmi::NeighborSearch ns(st.models[0], cell, 5.0);
    ns.populate(false);
    const std::vector<ModelNeighbour> near = AtomsNear(ns, st.models[0], p, 3.0);
    REQUIRE(near.size() == 2);
    CHECK(near[0].distance == Catch::Approx(0.0).margin(1e-6));
    CHECK(near[1].distance == Catch::Approx(std::sqrt(2.0) * 1.2).margin(1e-4));
    CHECK(near[1].pos.dist(gemmi::Position(near[1].image.apply(p))) == Catch::Approx(0.0).margin(1e-6));

    // and an atom across the face of the cell, a lattice translate away
    gemmi::Structure box = empty_box("P 1", 20.0);
    add_atom(box, "HOH", "O", gemmi::El::O, gemmi::Position(0.5, 10, 10));
    add_atom(box, "HOH", "O", gemmi::El::O, gemmi::Position(19.5, 10, 10));
    gemmi::UnitCell bcell = box.cell;
    bcell.set_cell_images_from_spacegroup(box.find_spacegroup());
    gemmi::NeighborSearch bns(box.models[0], bcell, 5.0);
    bns.populate(false);
    const std::vector<ModelNeighbour> across = AtomsNear(bns, box.models[0], gemmi::Position(0.5, 10, 10), 3.0);
    REQUIRE(across.size() == 2);
    CHECK(across[1].residue == 1);
    CHECK(across[1].distance == Catch::Approx(1.0).margin(1e-6));
    CHECK(across[1].pos.x == Catch::Approx(-0.5).margin(1e-6));
}

TEST_CASE("SolventSites_chemistry", "[Solvent]") {
    const gemmi::MonLib &monlib = EmbeddedMonomerLibrary();
    auto chem = [&](const std::string &res, const std::string &atom, gemmi::El el) {
        gemmi::Residue r;
        r.name = res;
        gemmi::Atom a;
        a.name = atom;
        a.element = gemmi::Element(el);
        r.atoms.push_back(a);
        return ChemistryOf(monlib, r, r.atoms[0], false);
    };
    const AtomChemistry nz = chem("LYS", "NZ", gemmi::El::N);
    CHECK((nz.donor && !nz.acceptor && nz.positive && !nz.negative));
    const AtomChemistry nh1 = chem("ARG", "NH1", gemmi::El::N);
    CHECK((nh1.donor && nh1.positive));
    const AtomChemistry od1 = chem("ASP", "OD1", gemmi::El::O);
    CHECK((od1.acceptor && !od1.donor && od1.negative));
    const AtomChemistry op1 = chem("G", "OP1", gemmi::El::O);
    CHECK((op1.acceptor && op1.negative));
    const AtomChemistry n7 = chem("G", "N7", gemmi::El::N);
    CHECK((n7.acceptor && !n7.donor && !n7.negative));
    const AtomChemistry n3 = chem("U", "N3", gemmi::El::N);
    CHECK((n3.donor && !n3.acceptor));
    const AtomChemistry o2p = chem("A", "O2'", gemmi::El::O);
    CHECK((o2p.donor && o2p.acceptor));
    const AtomChemistry carbonyl = chem("GLY", "O", gemmi::El::O);
    CHECK((carbonyl.acceptor && !carbonyl.donor && !carbonyl.negative));
    const AtomChemistry nd2 = chem("ASN", "ND2", gemmi::El::N);
    CHECK((nd2.flip_ambiguous && nd2.donor && nd2.acceptor));
    const AtomChemistry ring = chem("HIS", "NE2", gemmi::El::N);
    CHECK((ring.flip_ambiguous && ring.donor && ring.acceptor));
    const AtomChemistry water = chem("HOH", "O", gemmi::El::O);
    CHECK((water.donor && water.acceptor));
    CHECK(!chem("GLY", "CA", gemmi::El::C).acceptor);
}

TEST_CASE("SolventSites_octahedral_Mg", "[Solvent]") {
    // Mg(H2O)6, Mg-O 2.07 A
    gemmi::Structure st = empty_box();
    add_atom(st, "MG", "MG", gemmi::El::Mg, CENTRE);
    for (int axis = 0; axis < 3; axis++)
        for (double sign : {-1.0, 1.0}) {
            gemmi::Position p = CENTRE;
            (axis == 0 ? p.x : axis == 1 ? p.y : p.z) += sign * 2.07;
            add_atom(st, "HOH", "O", gemmi::El::O, p);
        }
    const std::vector<SolventSite> sites = AnalyseSolventSites(st, 0.0, SolventMaps{});
    REQUIRE(sites.size() == 7);
    const SolventSite &mg = site_labelled(sites, "MG");
    CHECK(mg.n_19_23 == 6);
    CHECK(mg.n_23_26 + mg.n_26_32 + mg.n_32_34 == 0);
    CHECK(mg.bvs[0] == Catch::Approx(6 * std::exp((1.69 - 2.07) / 0.37)).epsilon(1e-6));
    CHECK(mg.nvecsum[0] == Catch::Approx(0.0).margin(1e-9));
    CHECK(mg.min_ligand_distance == Catch::Approx(2.07 * std::sqrt(2.0)).epsilon(1e-6));
    CHECK(mg.best == SolventSpecies::Mg);
    CHECK(!mg.ambiguous);
    CHECK(SolventLabelAgrees(mg));
    // its waters are waters: each is bound to the Mg, and the other five are its cis and trans neighbours
    for (const SolventSite &s : sites)
        if (s.label == "HOH") {
            CHECK(s.nearest_positive == Catch::Approx(2.07).epsilon(1e-6));
            CHECK(s.best == SolventSpecies::Water);
        }
}

TEST_CASE("SolventSites_K_eight_carbonyls", "[Solvent]") {
    // K+ in a square antiprism of eight carbonyl O at 2.8 A
    gemmi::Structure st = empty_box();
    add_atom(st, "K", "K", gemmi::El::K, CENTRE);
    const double h = 1.3, r = std::sqrt(2.8 * 2.8 - h * h);
    for (int k = 0; k < 8; k++) {
        const double phi = (k < 4 ? k * 90.0 : k * 90.0 + 45.0) * 3.14159265358979 / 180.0;
        add_atom(st, "GLY", "O", gemmi::El::O,
                 gemmi::Position(CENTRE.x + r * std::cos(phi), CENTRE.y + r * std::sin(phi), CENTRE.z + (k < 4 ? h : -h)));
    }
    const std::vector<SolventSite> sites = AnalyseSolventSites(st, 0.0, SolventMaps{});
    REQUIRE(sites.size() == 1);
    const SolventSite &k = sites[0];
    CHECK(k.n_26_32 == 8);
    CHECK(k.acceptors == 8);
    CHECK(k.donors == 0);
    CHECK(k.bvs[2] == Catch::Approx(8 * std::exp((2.13 - 2.8) / 0.37)).epsilon(1e-6));
    CHECK(k.best == SolventSpecies::K);
    CHECK(!k.ambiguous);
    CHECK(k.score[static_cast<int>(SolventSpecies::K)] > k.score[static_cast<int>(SolventSpecies::Water)] + 5);
}

TEST_CASE("SolventSites_Cl_four_donors", "[Solvent]") {
    // Cl- with four backbone N-H donors at 3.2 A, tetrahedral
    gemmi::Structure st = empty_box();
    add_atom(st, "CL", "CL", gemmi::El::Cl, CENTRE);
    const double t = 3.2 / std::sqrt(3.0);
    for (const auto &[x, y, z] : std::vector<std::array<double, 3>>{{1, 1, 1}, {1, -1, -1}, {-1, 1, -1}, {-1, -1, 1}})
        add_atom(st, "GLY", "N", gemmi::El::N, gemmi::Position(CENTRE.x + t * x, CENTRE.y + t * y, CENTRE.z + t * z));
    const std::vector<SolventSite> sites = AnalyseSolventSites(st, 0.0, SolventMaps{});
    REQUIRE(sites.size() == 1);
    const SolventSite &cl = sites[0];
    CHECK(cl.n_26_32 + cl.n_32_34 == 4);
    CHECK(cl.donors == 4);
    CHECK(cl.best == SolventSpecies::Cl);
    CHECK(!cl.ambiguous);
    // a water accepts two hydrogen bonds, not four
    CHECK(cl.score[static_cast<int>(SolventSpecies::Cl)] > cl.score[static_cast<int>(SolventSpecies::Water)] + 2);
}

TEST_CASE("SolventSites_tetrahedral_water", "[Solvent]") {
    // a water donating to two carbonyl O at 2.8 A and accepting from two N-H at 2.9 A, tetrahedral
    gemmi::Structure st = empty_box();
    add_atom(st, "HOH", "O", gemmi::El::O, CENTRE);
    const std::vector<std::array<double, 3>> dirs = {{1, 1, 1}, {1, -1, -1}, {-1, 1, -1}, {-1, -1, 1}};
    for (size_t i = 0; i < dirs.size(); i++) {
        const double d = (i < 2 ? 2.8 : 2.9) / std::sqrt(3.0);
        add_atom(st, "GLY", i < 2 ? "O" : "N", i < 2 ? gemmi::El::O : gemmi::El::N,
                 gemmi::Position(CENTRE.x + d * dirs[i][0], CENTRE.y + d * dirs[i][1], CENTRE.z + d * dirs[i][2]));
    }
    const std::vector<SolventSite> sites = AnalyseSolventSites(st, 0.0, SolventMaps{});
    REQUIRE(sites.size() == 1);
    const SolventSite &w = sites[0];
    CHECK(w.n_26_32 == 4);
    CHECK(w.donors == 2);
    CHECK(w.acceptors == 2);
    CHECK(w.best == SolventSpecies::Water);
    CHECK(!w.ambiguous);
    CHECK(SolventLabelAgrees(w));
}

TEST_CASE("SolventSites_coordinate_error_blurs_the_bands", "[Solvent]") {
    // the octahedral Mg again, its B high and the resolution poor: the bands of Mg and Na overlap, so the score
    // can no longer prefer Mg by much over Na - but still not water, whose shell is far too far out
    gemmi::Structure st = empty_box();
    add_atom(st, "MG", "MG", gemmi::El::Mg, CENTRE);
    for (int axis = 0; axis < 3; axis++)
        for (double sign : {-1.0, 1.0}) {
            gemmi::Position p = CENTRE;
            (axis == 0 ? p.x : axis == 1 ? p.y : p.z) += sign * 2.07;
            add_atom(st, "HOH", "O", gemmi::El::O, p);
        }
    for (gemmi::Residue &r : st.models[0].chains[0].residues)
        r.atoms[0].b_iso = 80.0f;
    const SolventSite sharp = site_labelled(AnalyseSolventSites(st, 1.5, SolventMaps{}), "MG");
    const SolventSite blurred = site_labelled(AnalyseSolventSites(st, 3.5, SolventMaps{}), "MG");
    CHECK(blurred.sigma_xyz > sharp.sigma_xyz);
    const int mg = static_cast<int>(SolventSpecies::Mg), na = static_cast<int>(SolventSpecies::Na);
    CHECK(blurred.score[mg] - blurred.score[na] < sharp.score[mg] - sharp.score[na]);
    CHECK(blurred.best != SolventSpecies::Water);
}

namespace {

// A residue of the embedded dictionary at its ideal coordinates (hydrogens left out), moved by `shift`.
void add_residue(gemmi::Structure &st, const std::string &code, const gemmi::Position &shift) {
    gemmi::Chain &ch = st.models[0].chains[0];
    const gemmi::ChemComp &cc = EmbeddedMonomerLibrary().monomers.at(code);
    gemmi::Residue r;
    r.name = code;
    r.seqid = gemmi::SeqId(static_cast<int>(ch.residues.size()) + 1, ' ');
    for (const gemmi::ChemComp::Atom &ca : cc.atoms) {
        if (ca.is_hydrogen())
            continue;
        gemmi::Atom a;
        a.name = ca.id;
        a.element = ca.el;
        a.pos = gemmi::Position(ca.xyz) + shift;
        a.occ = 1.0f;
        a.b_iso = 30.0f;
        r.atoms.push_back(a);
    }
    ch.residues.push_back(r);
}

// The model's density on a grid (no blur beyond the atoms' B), as a stand-in for the 2mFo-DFc map.
gemmi::Grid<float> model_map(const gemmi::Structure &st, double d_min) {
    gemmi::DensityCalculator<gemmi::IT92<float>, float> dc;
    dc.d_min = d_min;
    dc.rate = 2.0;
    dc.set_grid_cell_and_spacegroup(st);
    dc.put_model_density_on_grid(st.models[0]);
    return dc.grid;
}

}  // namespace

// The electrons a site shows, its own density fitted with its neighbours' held: a Cl- among four donors reads ~18
// and is called Cl whichever label it carries; a water there reads as a 10-electron scatterer and is called water.
TEST_CASE("SolventSites_electrons_decide_halide_from_water", "[Solvent]") {
    for (const bool chloride : {true, false}) {
        CAPTURE(chloride);
        gemmi::Structure st = empty_box();
        // the macromolecule the map's scale is read from
        for (int k = 0; k < 6; k++)
            add_residue(st, "TRP", gemmi::Position(5 + 5 * k, k % 2 == 0 ? 6 : 34, k % 3 == 0 ? 8 : 30));
        const double t = 3.2 / std::sqrt(3.0);
        for (const auto &[x, y, z] : std::vector<std::array<double, 3>>{{1, 1, 1}, {1, -1, -1}, {-1, 1, -1}, {-1, -1, 1}})
            add_atom(st, "GLY", "N", gemmi::El::N, gemmi::Position(CENTRE.x + t * x, CENTRE.y + t * y, CENTRE.z + t * z));
        if (chloride) {
            add_atom(st, "CL", "CL", gemmi::El::Cl, CENTRE);
            st.models[0].chains[0].residues.back().atoms[0].charge = -1;
        } else
            add_atom(st, "HOH", "O", gemmi::El::O, CENTRE);
        for (gemmi::Residue &r : st.models[0].chains[0].residues)
            for (gemmi::Atom &a : r.atoms)
                a.b_iso = 30.0f;
        const gemmi::Grid<float> map = model_map(st, 1.5);
        // labelled water either way
        gemmi::Residue &site = st.models[0].chains[0].residues.back();
        site.name = "HOH";
        site.atoms[0] = gemmi::Atom();
        site.atoms[0].name = "O";
        site.atoms[0].element = gemmi::El::O;
        site.atoms[0].pos = CENTRE;
        site.atoms[0].occ = 1.0f;
        site.atoms[0].b_iso = 30.0f;
        SolventMaps maps;
        maps.map = &map;
        maps.map_rms = 1.0;
        const SolventSite &s = site_labelled(AnalyseSolventSites(st, 1.5, maps), "HOH");
        CHECK(s.electrons == Catch::Approx(chloride ? 18.0 : 8.0).epsilon(0.25));   // the model's water: O only
        CHECK(s.best == (chloride ? SolventSpecies::Cl : SolventSpecies::Water));
    }
}

// A divalent ion bound to a purine N7 lies in the base's plane, 3.5-3.9 A from its O6 (NAR 2017): one 1.5 A out of the
// plane is flagged, and costs Mg - not Na - its score.
TEST_CASE("SolventSites_N7_out_of_plane", "[Solvent]") {
    auto site_at = [](bool in_plane) {
        gemmi::Structure st = empty_box();
        add_residue(st, "G", CENTRE);
        const gemmi::Residue &g = st.models[0].chains[0].residues[0];
        const gemmi::Position n7 = g.find_atom("N7", '*')->pos, c5 = g.find_atom("C5", '*')->pos,
                              c8 = g.find_atom("C8", '*')->pos;
        // outward from N7, 10 degrees off the bisector of C5-N7-C8 away from O6 (3.8 A from it), in the plane; or
        // tilted 46 degrees out of it; 2.2 A from N7 either way
        gemmi::Vec3 out = (n7 - c5).normalized() + (n7 - c8).normalized();
        out = out.normalized();
        const gemmi::Vec3 normal = (c5 - n7).cross(c8 - n7).normalized();
        const double turn = -10.0 * 3.14159265358979 / 180.0;
        const gemmi::Vec3 in_dir = out * std::cos(turn) + normal.cross(out) * std::sin(turn);
        const gemmi::Vec3 dir = in_plane ? in_dir : (in_dir * std::cos(0.8) + normal * std::sin(0.8));
        add_atom(st, "MG", "MG", gemmi::El::Mg, n7 + gemmi::Position(dir * 2.2));
        return site_labelled(AnalyseSolventSites(st, 0.0, SolventMaps{}), "MG");
    };
    const SolventSite in = site_at(true), off = site_at(false);
    CHECK(in.n7_ligands == 1);
    CHECK(in.n7_off_plane == 0);
    CHECK(off.n7_ligands == 1);
    CHECK(off.n7_off_plane == 1);
    const int mg = static_cast<int>(SolventSpecies::Mg), na = static_cast<int>(SolventSpecies::Na);
    CHECK(off.score[mg] - off.score[na] < in.score[mg] - in.score[na]);
}