neta.cpp

// SPDX-License-Identifier: GPL-3.0-or-later
// Copyright (c) 2026 Team Dissolve and contributors

#include "neta/neta.h"
#include "classes/species.h"
#include "data/elements.h"
#include "data/ff/ff.h"
#include "templates/algorithms.h"
#include <format>
#include <gtest/gtest.h>

namespace UnitTest
{
class NETATest : public ::testing::Test
{
    public:
    NETATest()
    {
        methane_.load("species/methane.toml");
        ethane_.load("species/ethane.toml");
        rings_.load("species/rings.toml");
        difluorobenzene_.load("species/difluorobenzene.toml");
        geometric_.load("species/geometric.toml");
        methanol_.load("species/methanol.toml");
    }

    protected:
    Species methane_, ethane_, rings_, difluorobenzene_, geometric_, methanol_;

    protected:
    // Test NETA description on all atom in molecule, expecting success for supplied atom indices
    void testNETA(std::string_view title, const Species &sp, const NETADefinition &neta,
                  const std::vector<int> &matchingIndices)
    {
        std::cout << std::format("Testing: {}...", title) << std::endl;
        auto indexString = joinStrings(matchingIndices, ", ");
        std::cout << std::format("-- Species '{}', expected matching atoms: ({})", sp.name(), indexString) << std::endl;

        for (const auto &i : sp.atoms())
        {
            auto score = neta.score(&i);
            if (score != -1)
                std::cout << std::format("-- NETA score for atom {} is {}.", i.index(), score) << std::endl;
            if (std::find(matchingIndices.begin(), matchingIndices.end(), i.index()) != matchingIndices.end())
                EXPECT_NE(score, NETANode::NoMatch);
            else
                EXPECT_EQ(score, NETANode::NoMatch);
        }
    }
    // Test NETA description on specific atom in molecule, expecting the match path to contain the supplied atom indices
    NETAMatchedGroup testNETAMatchPath(std::string_view title, const Species &sp, const NETADefinition &neta,
                                       int targetAtomIndex, const std::vector<int> &matchingIndices)
    {
        std::cout << std::format("Path Test: {}, atom {}...", title, targetAtomIndex) << std::endl;
        auto indexString = joinStrings(matchingIndices, ", ");
        std::cout << std::format("-- Species '{}', expected matched atom set : {}", sp.name(), indexString) << std::endl;

        auto matchedPath = neta.matchedPath(&sp.atom(targetAtomIndex));
        std::cout << std::format("-- Actual matched atom set : {}",
                                 joinStrings(matchedPath.set(), " ", [](const auto *i) { return i->index(); }))
                  << std::endl;
        EXPECT_EQ(matchedPath.set().size(), matchingIndices.size());

        for (auto *i : matchedPath.set())
            EXPECT_TRUE(std::find(matchingIndices.begin(), matchingIndices.end(), i->index()) != matchingIndices.end());

        return matchedPath;
    }
    // Test identifier group specified
    void testIdentifiers(const NETAMatchedGroup &matchedPath, std::string idName, const std::vector<int> &matchingIndices)
    {
        std::cout << std::format("Identifier Test: {}", idName) << std::endl;
        auto indexString = joinStrings(matchingIndices, ", ");

        std::cout << std::format("-- Expected identified atoms : {}", indexString) << std::endl;
        auto &ids = matchedPath.identifiers();
        auto it = ids.find(idName);
        EXPECT_TRUE(it != ids.end());
        if (it == ids.end())
        {
            std::cout << std::format("Identifier group '{}' not present in matched data.", idName) << std::endl;
            return;
        }

        auto &group = it->second;

        std::cout << std::format("-- Actual identified atoms : {}",
                                 joinStrings(group, " ", [](const auto *i) { return i->index(); }))
                  << std::endl;

        EXPECT_EQ(group.size(), matchingIndices.size());

        for (auto *i : group)
            EXPECT_TRUE(std::find(matchingIndices.begin(), matchingIndices.end(), i->index()) != matchingIndices.end());
    }
};

TEST_F(NETATest, Syntax)
{
    NETADefinition neta;

    // Empty NETA string
    EXPECT_TRUE(neta.create(""));
    // Simple Connectivity
    EXPECT_TRUE(neta.create("-H,-C,-O"));
    // Or'd Connectivity
    EXPECT_TRUE(neta.create("-[Sc,Ti,V,Cr,Mn,Fe,Co,Ni,Cu]"));
    // Brackets
    EXPECT_TRUE(neta.create("(-H,(-C,(-O)))"));
    // Nested Connectivity
    EXPECT_TRUE(neta.create("-H(-He(-Li(-Be(-B(-C)))))"));
    // Modifiers
    EXPECT_TRUE(neta.create("-C(n>=2,-H(n=3),-N(n<4,nbonds=1,root)),nbonds<8"));
    // Rings
    EXPECT_TRUE(neta.create("ring(size=6),-C(ring(size>=5,Sc,Ti,V,Cr,Mn))"));
    // Or'd Nodes
    EXPECT_TRUE(neta.create("-H|-C|ring()"));
    // Or'd Node Sequence
    EXPECT_TRUE(neta.create("-H|-C(nh=4),ring()|-Zn"));
    // Or'd Node Sequence with brackets
    EXPECT_TRUE(neta.create("-H|(-C(nh=4),ring()|(-Zn))"));
    // Or'd Node Sequence in bracketed part
    EXPECT_TRUE(neta.create("-V(-H|-C(nh=4),ring()|-Zn)"));
    // Error - Dot not Comma
    EXPECT_FALSE(neta.create("-H.-C,-O"));
    // Error - Double '-'
    EXPECT_FALSE(neta.create("-H,-C,--O"));
    // Error - Missing bracket 1
    EXPECT_FALSE(neta.create("-H,-C(-O"));
    // Error - Mis-placed bracket 2
    EXPECT_FALSE(neta.create("-H,-C,-O)"));
    // Error - Unknown element
    EXPECT_FALSE(neta.create("-C,-B,-A]"));
    // Error - Unknown contextual
    EXPECT_FALSE(neta.create("-C(nangles=5)"));
    // Error - Unknown operator
    EXPECT_FALSE(neta.create("-C(nbonds<>5)"));
    // Error - Bad syntax
    EXPECT_FALSE(neta.create("?O,"));
    // Error - Double ','
    EXPECT_FALSE(neta.create("-H,-C,,-O"));
}

TEST_F(NETATest, Matching)
{
    NETADefinition neta;

    EXPECT_TRUE(neta.create("-C"));
    testNETA("Any atom bound to carbon", methane_, neta, {1, 2, 3, 4});

    EXPECT_TRUE(neta.create("nbonds=4"));
    testNETA("Any atom with four bonds", methane_, neta, {0});

    EXPECT_TRUE(neta.create("-C"));
    testNETA("Any atom bound to carbon", ethane_, neta, {0, 1, 2, 3, 4, 5, 6, 7});

    EXPECT_TRUE(neta.create("nh=3"));
    testNETA("Any atom with three hydrogens attached", ethane_, neta, {0, 1});

    EXPECT_TRUE(neta.create("ring()"));
    testNETA("Atom in any ring", rings_, neta, {0, 1, 2, 3, 4, 5, 10, 11});

    EXPECT_TRUE(neta.create("ring(size=6)"));
    testNETA("Atom in six-membered ring", rings_, neta, {0, 1, 2, 3, 4, 5});

    EXPECT_TRUE(neta.create("ring(n=0)"));
    testNETA("Any atom not in a ring", rings_, neta, {6, 7, 8, 9, 12, 13, 14, 15, 16, 17});

    EXPECT_TRUE(neta.create("nbonds=1,-C(ring(size=6))"));
    testNETA("Hydrogen atom attached to atom within a six-membered ring", rings_, neta, {6, 7, 8, 9});

    EXPECT_TRUE(neta.create("ring(size=6), ring(size=4)"));
    testNETA("Atom at junction of four and six-membered ring", rings_, neta, {0, 1});

    EXPECT_TRUE(neta.create("-N, ring(size=6), ring(size=4)"));
    testNETA("Atom at junction of four and six-membered ring, adjacent to nitrogen", rings_, neta, {1});

    EXPECT_TRUE(neta.create("-N,ring(n=0)"));
    testNETA("Atom adjacent to nitrogen, but not in a ring ", rings_, neta, {12});

    EXPECT_TRUE(neta.create("-N,ring(n=0) | -N,ring(size=6), ring(size=4)"));
    testNETA("Either of the previous atoms", rings_, neta, {1, 12});

    EXPECT_TRUE(neta.create("-C(nh=3),nbonds=1"));
    testNETA("Hydrogen atoms present in CH3 group", rings_, neta, {13, 14, 15});

    EXPECT_TRUE(neta.create("?C,!ring(size=6)"));
    testNETA("Any carbon except one in a 6-membered ring", rings_, neta, {11, 12});

    EXPECT_TRUE(neta.create("?C,!(nh=2 | -C(-N))"));
    testNETA("Any carbon except one with two hydrogens or which is two bonds away from a nitrogen", rings_, neta,
             {1, 3, 4, 5, 12});
}

TEST_F(NETATest, Creation)
{
    // Carbon atom - full description
    NETADefinition neta(&methane_.atom(0));
    EXPECT_EQ(neta.definitionString(), "nbonds=4,nh=4");
    testNETA("Carbon atom in methane", methane_, neta, {0});

    // Methane hydrogen atom
    // -- Basic connectivity
    neta.create(&methane_.atom(1), 0);
    EXPECT_EQ(neta.definitionString(), "nbonds=1,-C");
    testNETA("Hydrogen atom in methane", methane_, neta, {1, 2, 3, 4});
    // -- Primary neighbour connectivity
    neta.create(&methane_.atom(1), 1);
    EXPECT_EQ(neta.definitionString(), "nbonds=1,-C(nbonds=4,nh=4)");
    testNETA("Hydrogen atom in methane", methane_, neta, {1, 2, 3, 4});
    // -- Secondary neighbour connectivity (equivalent to primary)
    neta.create(&methane_.atom(1), 2);
    EXPECT_EQ(neta.definitionString(), "nbonds=1,-C(nbonds=4,nh=4)");
    testNETA("Hydrogen atom in methane", methane_, neta, {1, 2, 3, 4});
    // -- Basic connectivity including the root element
    neta.create(&methane_.atom(1), 0,
                Flags<NETADefinition::NETACreationFlags>(NETADefinition::NETACreationFlags::IncludeRootElement));
    EXPECT_EQ(neta.definitionString(), "?H, nbonds=1,-C");
    testNETA("Hydrogen atom in methane", methane_, neta, {1, 2, 3, 4});
    neta.create(&methane_.atom(1), 1,
                Flags<NETADefinition::NETACreationFlags>(NETADefinition::NETACreationFlags::IncludeRootElement));
    EXPECT_EQ(neta.definitionString(), "?H, nbonds=1,-C(nbonds=4,nh=4)");
}

TEST_F(NETATest, Forcefield)
{
    class Forcefield_TEST : public Forcefield
    {
        protected:
        bool setUp() override
        {
            addAtomType(Elements::C, 1, "C6", "ring(size=6)", "Carbon in 6-membered ring", 0.0, {0.0, 0.0});
            addAtomType(Elements::C, 2, "C4", "ring(size=4),nh=2", "Carbon in 4-membered ring", 0.0, {0.0, 0.0});
            addAtomType(Elements::N, 3, "N4", "ring(size=4)", "Nitrogen in 4-membered ring", 0.0, {0.0, 0.0});

            return true;
        }
        std::string_view name() const override { return "Test Forcefield for NETA."; }
        std::string_view description() const override { return "Test Forcefield for NETA."; }
        ShortRangeFunctions::Form shortRangeForm() const override { return ShortRangeFunctions::Form::LennardJones; }
    };

    Forcefield_TEST testFF;
    testFF.prepare();
    testFF.createNETADefinitions();

    NETADefinition neta;

    EXPECT_TRUE(neta.create("-&C6,nbonds=1", &testFF));
    testNETA("Reference to type by name", rings_, neta, {6, 7, 8, 9});

    EXPECT_TRUE(neta.create("-&1,nbonds=1", &testFF));
    testNETA("Reference to type by ID", rings_, neta, {6, 7, 8, 9});

    std::cout << std::format("Testing: Non-existent type name...") << std::endl;
    EXPECT_FALSE(neta.create("-&C5", &testFF));

    std::cout << std::format("Testing: Non-existent type ID...") << std::endl;
    EXPECT_FALSE(neta.create("-&8", &testFF));
}

TEST_F(NETATest, Geometry)
{
    NETADefinition neta;

    neta.create("geometry=unbound");
    testNETA("Geometry = unbound", geometric_, neta, {0});

    neta.create("geometry=terminal");
    testNETA("Geometry = terminal [single bond]", geometric_, neta,
             {1, 2, 4, 5, 7, 8, 9, 11, 12, 13, 14, 16, 17, 18, 19, 20, 21, 23, 24, 25, 26, 27, 29, 30, 31, 33, 34, 35, 36});

    neta.create("geometry!=terminal");
    testNETA("Geometry != terminal [single bond]", geometric_, neta, {0, 3, 6, 10, 15, 22, 28, 32});

    neta.create("geometry=linear");
    testNETA("Geometry = linear", geometric_, neta, {3});

    neta.create("geometry=ts");
    testNETA("Geometry = TShape", geometric_, neta, {28});

    neta.create("geometry=tp");
    testNETA("Geometry = trigonal planar", geometric_, neta, {6});

    neta.create("geometry=tet");
    testNETA("Geometry = tetrahedral", geometric_, neta, {10});

    neta.create("geometry=sqp");
    testNETA("Geometry = square planar", geometric_, neta, {32});

    neta.create("geometry=tbp");
    testNETA("Geometry = trigonal bipyramidal", geometric_, neta, {22});

    neta.create("geometry=oct");
    testNETA("Geometry = octahedral", geometric_, neta, {15});
}

TEST_F(NETATest, FragmentMatching)
{
    NETADefinition neta;

    EXPECT_TRUE(neta.create("?C,-H(n=3)"));
    testNETAMatchPath("Methyl group C(0)", ethane_, neta, 0, {0, 2, 3, 4});

    EXPECT_TRUE(neta.create("?C,-C(-H(n=3))"));
    testNETAMatchPath("C(1) methyl plus C(0)", ethane_, neta, 0, {0, 1, 5, 6, 7});

    EXPECT_TRUE(neta.create("ring(size=4,C,C,C,N)"));
    testNETA("Any atom in a four-membered ring (explicit ring definition)", rings_, neta, {0, 1, 10, 11});

    EXPECT_TRUE(neta.create("ring(size=4,C(n=3),N)"));
    testNETA("Any atom in a four-membered ring (shortest ring definition)", rings_, neta, {0, 1, 10, 11});

    EXPECT_TRUE(neta.create("ring(size=4,C,C(n=2),N)"));
    testNETA("Any atom in a four-membered ring (unnecessary ring definition)", rings_, neta, {0, 1, 10, 11});

    EXPECT_TRUE(neta.create("ring(size=4,N)"));
    testNETA("Any atom in a four-membered ring (only one atom specified)", rings_, neta, {0, 1, 10, 11});

    EXPECT_TRUE(neta.create("ring(size=4,C(n=2))"));
    testNETA("Any atom in a four-membered ring (only one atom specified)", rings_, neta, {0, 1, 10, 11});

    EXPECT_TRUE(neta.create("ring(size=8,C(n=4),N,C(n=3))"));
    testNETA("Any atom in an eight-membered ring (split definition)", rings_, neta, {0, 1, 2, 3, 4, 5, 10, 11});

    EXPECT_TRUE(neta.create("ring(size=4,N(n=2),C(n=2))"));
    testNETA("No atoms - too many of atom type requested", rings_, neta, {});

    EXPECT_TRUE(neta.create("ring(size=8,C(n=4),N,C(n=4))"));
    testNETA("No atoms - too many for ring (9 vs 8)", rings_, neta, {});

    EXPECT_TRUE(neta.create("ring(size=6,C(n=8))"));
    testNETA("No atoms - too many for ring (8 vs 6)", rings_, neta, {});
}

TEST_F(NETATest, IdentifierMatching)
{
    NETADefinition neta;

    // Methanol identifying cog at the oxygen, x on the carbon, and y on the hydroxyl hydrogen
    EXPECT_TRUE(neta.create("?O,#cog,-C(#x),-H(-O(root),#y)"));
    auto matchedPath = testNETAMatchPath("Carbon and hydroxyl", methanol_, neta, 2, {0, 2, 5});
    for (auto &&[key, ids] : matchedPath.identifiers())
        std::cout << std::format("ID '{}' : {}", key, joinStrings(ids, " ", [](const auto *i) { return i->index(); }))
                  << std::endl;
    testIdentifiers(matchedPath, "cog", {2});
    testIdentifiers(matchedPath, "x", {0});
    testIdentifiers(matchedPath, "y", {5});

    // Difluorobenzene, identify carbon atoms attached to fluorines as 'cog'
    EXPECT_TRUE(neta.create("?C,ring(size=6,C(-H),C(-H),C(#cog,-F),C(-H),C(-H),C(#cog,-F))"));
    matchedPath = testNETAMatchPath("Whole molecule", difluorobenzene_, neta, 0, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11});
    testIdentifiers(matchedPath, "cog", {2, 5});

    // Difluorobenzene with full axis definition
    EXPECT_TRUE(neta.create("?C,ring(size=6,C(-H),C(-H),C(#cog,-F),C(-H),C(#y,-H),C(#[cog,x],-F))"));
    matchedPath = testNETAMatchPath("Whole molecule", difluorobenzene_, neta, 0, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11});
    testIdentifiers(matchedPath, "cog", {2, 5});
    testIdentifiers(matchedPath, "x", {5});
    testIdentifiers(matchedPath, "y", {4});

    // Difluorobenzene with full axis definition (separate cog and x specification for one carbon
    EXPECT_TRUE(neta.create("?C,ring(size=6,C(-H),C(-H),C(#cog,-F),C(-H),C(#y,-H),C(#cog,#x,-F))"));
    matchedPath = testNETAMatchPath("Whole molecule", difluorobenzene_, neta, 0, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11});
    testIdentifiers(matchedPath, "cog", {2, 5});
    testIdentifiers(matchedPath, "x", {5});
    testIdentifiers(matchedPath, "y", {4});
}

} // namespace UnitTest