// 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 <cmath>
#include <string>
#include <vector>

#include "../image_analysis/structure_refinement/SolventBuild.h"

namespace {

constexpr double EDGE = 24.0;   // A, a P1 box

gemmi::Structure box() {
    gemmi::Structure st;
    st.cell = gemmi::UnitCell(EDGE, EDGE, EDGE, 90, 90, 90);
    st.spacegroup_hm = "P 1";
    st.setup_cell_images();
    st.models.emplace_back(1);
    st.models[0].chains.emplace_back("A");
    return st;
}

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);
}

// A map of the box, 0.3 A grid, in units of its rms (1), with a Gaussian peak of `height` at each of `peaks`.
gemmi::Grid<float> map_with(const gemmi::UnitCell &cell, const std::vector<gemmi::Position> &peaks, float height,
                            float base = 0.0f) {
    gemmi::Grid<float> g;
    g.unit_cell = cell;
    g.spacegroup = gemmi::find_spacegroup_by_name("P 1");
    g.set_size(80, 80, 80);
    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 gemmi::Position p = g.get_position(u, v, w);
                float value = base;
                for (const gemmi::Position &q : peaks) {
                    const double d = cell.find_nearest_image(p, q, gemmi::Asu::Any).dist();
                    value += static_cast<float>(height * std::exp(-d * d / (2 * 0.35 * 0.35)));
                }
                g.data[g.index_q(u, v, w)] = value;
            }
    return g;
}

int count_kind(const SolventRound &r, SolventEdit::Kind kind) {
    int n = 0;
    for (const SolventEdit &e : r.edits)
        n += e.kind == kind;
    return n;
}

}  // namespace

// A difference peak 2.8 A from a water becomes a water; one with only a carbon 3.3 A away (no hydrogen-bond partner)
// does not, and is not reported either; the new water goes into the model as HOH.
TEST_CASE("SolventBuild_peak_with_and_without_partner", "[Solvent]") {
    gemmi::Structure st = box();
    add_atom(st, "HOH", "O", gemmi::El::O, {6, 6, 6});
    add_atom(st, "UNL", "C1", gemmi::El::C, {16, 16, 16});
    const gemmi::Position with_partner(8.8, 6, 6), without(16, 16, 19.3);
    const gemmi::Grid<float> diff = map_with(st.cell, {with_partner, without}, 6.0f);
    const gemmi::Grid<float> map = map_with(st.cell, {}, 0.0f, 1.5f);
    const SolventMaps maps{&map, 1.0, &diff, 1.0, nullptr, 0};

    SolventRound r = PickSolvent(st, maps, {}, {}, {}, 300, 1);
    REQUIRE(r.added == 1);
    CHECK(r.not_solvent == 0);
    CHECK(r.refused["no hydrogen-bond partner"] == 1);
    const SolventEdit &e = r.edits.at(0);
    CHECK(e.kind == SolventEdit::Kind::Added);
    CHECK(st.cell.find_nearest_image(e.pos, with_partner, gemmi::Asu::Any).dist() < 0.25);
    CHECK(e.difference_sigma == Catch::Approx(6.0).margin(0.5));

    ApplySolventEdits(st, r.edits);
    const gemmi::Residue &added = st.models[0].chains[0].residues.back();
    CHECK(added.name == "HOH");
    CHECK(added.seqid.num == 3);
    CHECK(added.atoms.at(0).b_iso == Catch::Approx(20.0));   // its partner's
    CHECK(r.edits.at(0).site == "HOH A3");
}

// A sulfate-like cluster - a peak with four more at 1.5 A on a tetrahedron - beside a water is no water: the centre
// is reported as a tetrahedral group, its corners as atoms of one molecule; nothing is added.
TEST_CASE("SolventBuild_tetrahedral_cluster_refused", "[Solvent]") {
    gemmi::Structure st = box();
    add_atom(st, "HOH", "O", gemmi::El::O, {6, 6, 6});
    const gemmi::Position centre(9.2, 6, 6);
    const double t = 1.5 / std::sqrt(3.0);
    std::vector<gemmi::Position> peaks = {centre};
    for (const gemmi::Position &d : {gemmi::Position(t, t, t), gemmi::Position(t, -t, -t), gemmi::Position(-t, t, -t),
                                     gemmi::Position(-t, -t, t)})
        peaks.push_back(centre + d);
    const gemmi::Grid<float> diff = map_with(st.cell, peaks, 6.0f);
    const gemmi::Grid<float> map = map_with(st.cell, {}, 0.0f, 1.5f);
    const SolventMaps maps{&map, 1.0, &diff, 1.0, nullptr, 0};

    const SolventRound r = PickSolvent(st, maps, {}, {}, {}, 300, 1);
    CHECK(r.added == 0);
    CHECK(count_kind(r, SolventEdit::Kind::NotSolvent) >= 1);
    bool tetrahedral = false;
    for (const SolventEdit &e : r.edits)
        if (e.kind == SolventEdit::Kind::NotSolvent && st.cell.find_nearest_image(e.pos, centre, gemmi::Asu::Any).dist() < 0.3)
            tetrahedral = e.reason.find("tetrahedral") != std::string::npos;
    CHECK(tetrahedral);
}

// A water in negative difference density is taken out; one without a partner only where the build added it.
TEST_CASE("SolventBuild_removal", "[Solvent]") {
    gemmi::Structure st = box();
    add_atom(st, "HOH", "O", gemmi::El::O, {6, 6, 6});
    add_atom(st, "HOH", "O", gemmi::El::O, {8.8, 6, 6});
    add_atom(st, "HOH", "O", gemmi::El::O, {16, 16, 16});
    const gemmi::Grid<float> diff = map_with(st.cell, {}, 0.0f);
    const gemmi::Grid<float> map = map_with(st.cell, {}, 0.0f, 1.5f);
    const SolventMaps maps{&map, 1.0, &diff, 1.0, nullptr, 0};
    std::vector<SolventSite> sites(3);
    for (int k = 0; k < 3; k++) {
        sites[k].label = "HOH";
        sites[k].chain = "A";
        sites[k].seqid = std::to_string(k + 1);
        sites[k].pos = st.models[0].chains[0].residues[k].atoms[0].pos;
        sites[k].map_sigma = 0.5;   // weak, which alone takes nothing out
        sites[k].n_26_32 = k < 2;
        sites[k].nearest_ion = INFINITY;
        sites[k].b_iso = 20;
    }
    sites[1].difference_sigma = -4.0;

    CHECK(PickSolvent(st, maps, sites, {}, {}, 300, 1).removed == 1);
    SolventRound r = PickSolvent(st, maps, sites, {"HOH A3"}, {}, 300, 1);
    CHECK(r.removed == 2);
    ApplySolventEdits(st, r.edits);
    REQUIRE(st.models[0].chains[0].residues.size() == 1);
    CHECK(st.models[0].chains[0].residues[0].seqid.num == 1);
}

// A water among six waters at 2.07 A relabelled Mg: the residue, element and atom name change and it gets a metal
// record for each of its six ligands; relabelled water again, the records go.
TEST_CASE("SolventBuild_relabel", "[Solvent]") {
    gemmi::Structure st = box();
    const gemmi::Position c(12, 12, 12);
    add_atom(st, "HOH", "O", gemmi::El::O, c);
    for (int axis = 0; axis < 3; axis++)
        for (double sign : {-1.0, 1.0}) {
            gemmi::Position p = c;
            (axis == 0 ? p.x : axis == 1 ? p.y : p.z) += sign * 2.07;
            add_atom(st, "HOH", "O", gemmi::El::O, p);
        }
    std::vector<SolventEdit> edits(1);
    edits[0].kind = SolventEdit::Kind::Relabelled;
    edits[0].site = "HOH A1";
    edits[0].chain = "A";
    edits[0].seqid = gemmi::SeqId(1, ' ');
    edits[0].species = "MG";
    ApplySolventEdits(st, edits);
    const gemmi::Residue &r = st.models[0].chains[0].residues[0];
    CHECK(r.name == "MG");
    CHECK(r.atoms[0].name == "MG");
    CHECK(r.atoms[0].element == gemmi::El::Mg);
    REQUIRE(st.connections.size() == 6);
    for (const gemmi::Connection &conn : st.connections)
        CHECK(conn.partner2.res_id.name == "MG");

    edits[0].site = "MG A1";
    edits[0].species = "HOH";
    ApplySolventEdits(st, edits);
    CHECK(st.models[0].chains[0].residues[0].name == "HOH");
    CHECK(st.models[0].chains[0].residues[0].atoms[0].element == gemmi::El::O);
    CHECK(st.connections.empty());
}