energy.h

#pragma once

#include "bonds.h"
#include "externalpotential.h" // Energybase implemented here
#include "potentials_base.h"
#include "sasa.h"
#include "space.h"
#include "aux/eigensupport.h"
#include "aux/pairmatrix.h"
#include "smart_montecarlo.h"
#include <range/v3/range/conversion.hpp>
#include <spdlog/spdlog.h>
#include <numeric>
#include <algorithm>
#include <concepts>

struct freesasa_parameters_fwd; // workaround for freesasa unnamed struct that cannot be forward
                                // declared

#if defined(__cpp_lib_parallel_algorithm) && __has_include(<tbb/tbb.h>)
#include <execution>
#endif

#if defined(__cpp_lib_parallel_algorithm) &&                                                       \
    __has_include(<tbb/tbb.h>) && ((defined(__clang__) && __clang_major__ >= 10) || (defined(__GNUC__) && __GNUC__ >= 10))
#define HAS_PARALLEL_TRANSFORM_REDUCE
#endif

namespace Faunus {

namespace ReactionCoordinate {
class ReactionCoordinateBase;
}

namespace pairpotential {
class PairPotential;
}

namespace Energy {

class Hamiltonian;

class ContainerOverlap : public EnergyTerm
{
  private:
    const Space& spc;
    bool groupIsOutsideContainer(const Change::GroupChange& group_change) const;
    double energyOfAllGroups() const;

  public:
    explicit ContainerOverlap(const Space& spc);
    double energy(const Change& change) override;
};

struct EwaldData
{
    using Tcomplex = std::complex<double>;
    Eigen::Matrix3Xd k_vectors; 
    Eigen::VectorXd Aks;        
    Eigen::VectorXcd Q_ion, Q_dipole;       
    double r_cutoff = 0;                    
    double n_cutoff = 0;                    
    double surface_dielectric_constant = 0; 
    double bjerrum_length = 0;              
    double kappa = 0;                       
    double kappa_squared = 0;               
    double alpha = 0;
    double const_inf = 0;
    double check_k2_zero = 0;
    bool use_spherical_sum = true;
    int num_kvectors = 0;
    Point box_length = {0.0, 0.0, 0.0}; 

    enum Policies
    {
        PBC,
        PBCEigen,
        IPBC,
        IPBCEigen,
        INVALID
    }; 

    Policies policy = PBC;             
    explicit EwaldData(const json& j); 
};

NLOHMANN_JSON_SERIALIZE_ENUM(EwaldData::Policies, {
                                                      {EwaldData::INVALID, nullptr},
                                                      {EwaldData::PBC, "PBC"},
                                                      {EwaldData::PBCEigen, "PBCEigen"},
                                                      {EwaldData::IPBC, "IPBC"},
                                                      {EwaldData::IPBCEigen, "IPBCEigen"},
                                                  })

void to_json(json& j, const EwaldData& d);

class EwaldPolicyBase
{
  public:
    std::string cite; 
    virtual ~EwaldPolicyBase() = default;
    virtual void
    updateBox(EwaldData&,
              const Point&) const = 0; 
    virtual void
    updateComplex(EwaldData& d,
                  const Space::GroupVector& groups) const = 0; 
    virtual void updateComplex(EwaldData& d, const Change& change, const Space::GroupVector& groups,
                               const Space::GroupVector& oldgroups)
        const = 0; 
    virtual double
    selfEnergy(const EwaldData& d, Change& change,
               Space::GroupVector& groups) = 0; 
    virtual double surfaceEnergy(
        const EwaldData& d, const Change& change,
        const Space::GroupVector& groups) = 0; 
    virtual double reciprocalEnergy(const EwaldData& d) = 0; 

    static auto mapGroupsToEigen(Space::GroupVector& groups)
    {
        auto is_partially_inactive = [](const Group& group) {
            return group.size() != group.capacity();
        };
        if (std::ranges::any_of(groups, is_partially_inactive)) {
            throw std::runtime_error("Eigen optimized Ewald not available with inactive groups");
        }
        auto first_particle = groups.front().begin();
        auto last_particle = groups.back().end();
        auto pos = asEigenMatrix(first_particle, last_particle,
                                 &Particle::pos); // N x 3
        auto charge = asEigenVector(first_particle, last_particle,
                                    &Particle::charge); // N x 1
        return std::make_tuple(pos, charge);
    }

    static auto mapGroupsToEigen(const Space::GroupVector& groups)
    {
        auto is_partially_inactive = [](const Group& group) {
            return group.size() != group.capacity();
        };
        if (std::ranges::any_of(groups, is_partially_inactive)) {
            throw std::runtime_error("Eigen optimized Ewald not available with inactive groups");
        }
        auto first_particle = groups.front().begin();
        auto last_particle = groups.back().end();
        auto pos = asEigenMatrix(first_particle, last_particle,
                                 &Particle::pos); // N x 3
        auto charge = asEigenVector(first_particle, last_particle,
                                    &Particle::charge); // N x 1
        return std::make_tuple(pos, charge);
    }

    static std::unique_ptr<EwaldPolicyBase> makePolicy(EwaldData::Policies); 
};

struct PolicyIonIon : public EwaldPolicyBase
{
    PolicyIonIon();
    void updateBox(EwaldData& d, const Point& box) const override;
    void updateComplex(EwaldData& data, const Space::GroupVector& groups) const override;
    void updateComplex(EwaldData& d, const Change& change, const Space::GroupVector& groups,
                       const Space::GroupVector& oldgroups) const override;
    double selfEnergy(const EwaldData& d, Change& change, Space::GroupVector& groups) override;
    double surfaceEnergy(const EwaldData& data, const Change& change,
                         const Space::GroupVector& groups) override;
    double reciprocalEnergy(const EwaldData& d) override;
};

struct PolicyIonIonEigen : public PolicyIonIon
{
    using PolicyIonIon::updateComplex;
    void updateComplex(EwaldData&, const Space::GroupVector&) const override;
    double reciprocalEnergy(const EwaldData&) override;
};

struct PolicyIonIonIPBC : public PolicyIonIon
{
    using PolicyIonIon::updateComplex;
    PolicyIonIonIPBC();
    void updateBox(EwaldData&, const Point&) const override;
    void updateComplex(EwaldData&, const Space::GroupVector&) const override;
    void updateComplex(EwaldData&, const Change&, const Space::GroupVector&,
                       const Space::GroupVector&) const override;
};

struct PolicyIonIonIPBCEigen : public PolicyIonIonIPBC
{
    using PolicyIonIonIPBC::updateComplex;
    void updateComplex(EwaldData&, const Space::GroupVector&) const override;
};

class Ewald : public EnergyTerm
{
  private:
    const Space& spc;
    EwaldData data;
    std::shared_ptr<EwaldPolicyBase> policy; 
    const Space::GroupVector* old_groups = nullptr;

  public:
    Ewald(const Space& spc, const EwaldData& data);
    Ewald(const json& j, const Space& spc);
    void init() override;
    void setOldGroups(const Space::GroupVector&
                          old_groups); 
    void updateState(const Change& change) override;
    double energy(const Change& change) override;
    void sync(EnergyTerm* energybase, const Change& change) override;
    void to_json(json& j) const override;
    void force(std::vector<Point>& forces) override; // update forces on all particles
};

class Isobaric : public EnergyTerm
{
  private:
    const Space& spc;
    double pressure = 0.0; 
    static const std::map<std::string, double>
        pressure_units; 
  public:
    Isobaric(const json& j, const Space& spc);
    double energy(const Change& change) override;
    void to_json(json& j) const override;
};

class Constrain : public EnergyTerm
{
  private:
    std::string type;
    std::unique_ptr<ReactionCoordinate::ReactionCoordinateBase> coordinate;
    std::optional<Faunus::pairpotential::HarmonicBond> harmonic;

  public:
    Constrain(const json& j, Space& space);
    double energy(const Change& change) override;
    void to_json(json& j) const override;
};

class Bonded : public EnergyTerm
{
  private:
    using BondVector = BasePointerVector<pairpotential::BondData>;
    const Space& spc;
    BondVector external_bonds;                
    std::map<int, BondVector> internal_bonds; 
    void updateGroupBonds(const Space::GroupType& group); 
    double sumBondEnergy(const BondVector& bonds) const;  
    double internalGroupEnergy(const Change::GroupChange& changed); 
    double sumEnergy(const BondVector& bonds,
                     const std::ranges::range auto& particle_indices) const;
    void updateInternalBonds(); 

  public:
    Bonded(const Space& spc, BondVector external_bonds);
    Bonded(const json& j, const Space& spc);
    void to_json(json& j) const override;
    double energy(const Change& change) override;    
    void force(std::vector<Point>& forces) override; 
};

double Bonded::sumEnergy(const Bonded::BondVector& bonds,
                         const std::ranges::range auto& particle_indices) const
{
    assert(std::is_sorted(particle_indices.begin(), particle_indices.end()));

    auto index_is_included = [&](auto index) {
        return std::binary_search(particle_indices.begin(), particle_indices.end(), index);
    };
    auto affected_bonds = bonds | std::views::filter([&](const auto& bond) {
                              return std::ranges::any_of(bond->indices, index_is_included);
                          });
    auto bond_energy = [dist = spc.geometry.getDistanceFunc()](const auto& bond) {
        return bond->energyFunc(dist);
    };
    return std::transform_reduce(affected_bonds.begin(), affected_bonds.end(), 0.0, std::plus<>(),
                                 bond_energy);
}

auto indexComplement(std::integral auto size, const std::ranges::range auto& range)
{
    using namespace std::views;
    return iota(0, static_cast<int>(size)) |
           filter([&](auto i) { return !std::binary_search(range.begin(), range.end(), i); });
}

template <pairpotential::RequirePairPotential TPairPotential,
          bool allow_anisotropic_pair_potential = true>
class PairEnergy
{
    const Space::GeometryType& geometry; 
    TPairPotential pair_potential;       
    Space& spc; 
    BasePointerVector<EnergyTerm>&
        potentials; 
  public:
    PairEnergy(Space& spc, BasePointerVector<EnergyTerm>& potentials)
        : geometry(spc.geometry)
        , spc(spc)
        , potentials(potentials)
    {
    }

    template <typename T> inline double potential(const T& a, const T& b) const
    {
        assert(&a != &b); // a and b cannot be the same particle
        if constexpr (allow_anisotropic_pair_potential) {
            const Point r = geometry.vdist(a.pos, b.pos);
            return pair_potential(a, b, r.squaredNorm(), r);
        }
        else {
            return pair_potential(a, b, geometry.sqdist(a.pos, b.pos), {0, 0, 0});
        }
    }

    // just a temporary placement until PairForce class template will be implemented
    template <typename ParticleType>
    inline Point force(const ParticleType& a, const ParticleType& b) const
    {
        assert(&a != &b); // a and b cannot be the same particle
        const Point b_towards_a = geometry.vdist(a.pos, b.pos); // vector b -> a = a - b
        return pair_potential.force(a, b, b_towards_a.squaredNorm(), b_towards_a);
    }

    template <typename... Args> inline auto operator()(Args&&... args)
    {
        return potential(std::forward<Args>(args)...);
    }

    void addPairPotentialSelfEnergy()
    {
        if (pair_potential.selfEnergy) { // only add if self energy is defined
            faunus_logger->debug("Adding self-energy from {} to hamiltonian", pair_potential.name);
            potentials.emplace_back<Energy::ParticleSelfEnergy>(spc, pair_potential.selfEnergy);
        }
    }

    void from_json(const json& j)
    {
        pairpotential::from_json(j, pair_potential);
        if (!pair_potential.isotropic && !allow_anisotropic_pair_potential) {
            throw std::logic_error("Only isotropic pair potentials are allowed.");
        }
        addPairPotentialSelfEnergy();
    }

    void to_json(json& j) const { pair_potential.to_json(j); }
};

template <typename T>
concept RequirePairEnergy = requires(T instance) {
    instance.potential(Particle(), Particle());
    instance.force(Particle(), Particle());
    instance.addPairPotentialSelfEnergy();
};

class EnergyAccumulatorBase
{
  protected:
    double value = 0.0;                                               
    using ParticleRef = const std::reference_wrapper<const Particle>; 
    using ParticlePair = std::pair<ParticleRef, ParticleRef>; 

  public:
    enum class Scheme
    {
        SERIAL,
        OPENMP,
        PARALLEL,
        INVALID
    };
    Scheme scheme = Scheme::SERIAL;

    EnergyAccumulatorBase(double value);
    virtual ~EnergyAccumulatorBase() = default;
    virtual void reserve(size_t number_of_particles);
    virtual void clear();
    virtual void from_json(const json& j);
    virtual void to_json(json& j) const;

    virtual explicit operator double();
    virtual EnergyAccumulatorBase& operator=(double new_value) = 0;
    virtual EnergyAccumulatorBase& operator+=(double new_value) = 0;
    virtual EnergyAccumulatorBase& operator+=(ParticlePair&& pair) = 0;

    template <typename TOtherAccumulator>
    inline EnergyAccumulatorBase& operator+=(TOtherAccumulator& acc)
    {
        value += static_cast<double>(acc);
        return *this;
    }
};

NLOHMANN_JSON_SERIALIZE_ENUM(EnergyAccumulatorBase::Scheme,
                             {{EnergyAccumulatorBase::Scheme::INVALID, nullptr},
                              {EnergyAccumulatorBase::Scheme::SERIAL, "serial"},
                              {EnergyAccumulatorBase::Scheme::OPENMP, "openmp"},
                              {EnergyAccumulatorBase::Scheme::PARALLEL, "parallel"}})

template <class T>
concept RequireEnergyAccumulator = std::is_base_of_v<EnergyAccumulatorBase, T>;

template <RequirePairEnergy PairEnergy>
class InstantEnergyAccumulator : public EnergyAccumulatorBase
{
  private:
    const PairEnergy&
        pair_energy; 

  public:
    InstantEnergyAccumulator(const PairEnergy& pair_energy, const double value = 0.0)
        : EnergyAccumulatorBase(value)
        , pair_energy(pair_energy)
    {
    }

    inline InstantEnergyAccumulator& operator=(const double new_value) override
    {
        value = new_value;
        return *this;
    }

    inline InstantEnergyAccumulator& operator+=(const double new_value) override
    {
        value += new_value;
        return *this;
    }

    inline InstantEnergyAccumulator& operator+=(ParticlePair&& pair) override
    {
        // keep this short to get inlined
        value += pair_energy.potential(pair.first.get(), pair.second.get());
        return *this;
    }

    void from_json(const json& j) override
    {
        EnergyAccumulatorBase::from_json(j);
        if (scheme != Scheme::SERIAL) {
            faunus_logger->warn("unsupported summation scheme; falling back to 'serial'");
        }
    }
};

template <RequirePairEnergy PairEnergy>
class DelayedEnergyAccumulator : public EnergyAccumulatorBase
{
  private:
    std::vector<ParticlePair> particle_pairs;
    const PairEnergy&
        pair_energy; 
    const size_t max_particles_in_buffer =
        10000; 

  public:
    explicit DelayedEnergyAccumulator(const PairEnergy& pair_energy, const double value = 0.0)
        : EnergyAccumulatorBase(value)
        , pair_energy(pair_energy)
    {
    }

    void reserve(size_t number_of_particles) override
    {
        try {
            number_of_particles = std::min(number_of_particles, max_particles_in_buffer);
            const auto number_of_pairs = (number_of_particles - 1U) * number_of_particles / 2U;
            faunus_logger->debug(
                std::format("reserving memory for {} energy pairs ({} MB)", number_of_pairs,
                            number_of_pairs * sizeof(ParticlePair) / (1024U * 1024U)));
            particle_pairs.reserve(number_of_pairs);
        }
        catch (std::exception& e) {
            throw std::runtime_error(std::format(
                "cannot allocate memory for energy pairs: {}. Use another summation policy.",
                e.what()));
        }
    }

    void clear() override
    {
        value = 0.0;
        particle_pairs.clear();
    }

    DelayedEnergyAccumulator& operator=(const double new_value) override
    {
        clear();
        value = new_value;
        return *this;
    }

    inline DelayedEnergyAccumulator& operator+=(const double new_value) override
    {
        value += new_value;
        return *this;
    }

    inline DelayedEnergyAccumulator& operator+=(ParticlePair&& pair) override
    {
        assert(particle_pairs.capacity() > 0);
        if (particle_pairs.size() == particle_pairs.capacity()) {
            operator double(); // sum stored pairs and reset buffer
        }
        particle_pairs.emplace_back(pair);
        return *this;
    }

    explicit operator double() override
    {
        switch (scheme) {
        case Scheme::OPENMP:
            value += accumulateOpenMP();
            break;
        case Scheme::PARALLEL:
            value += accumulateParallel();
            break;
        default:
            value += accumulateSerial();
        }
        particle_pairs.clear();
        return value;
    }

  private:
    double accumulateSerial() const
    {
        double sum = 0.0;
        for (const auto& [particle1, particle2] : particle_pairs) {
            sum += pair_energy.potential(particle1.get(), particle2.get());
        }
        return sum;
    }

    double accumulateParallel() const
    {
#if defined(HAS_PARALLEL_TRANSFORM_REDUCE)
        return std::transform_reduce(
            std::execution::par, particle_pairs.cbegin(), particle_pairs.cend(), 0.0, std::plus<>(),
            [&](const auto& pair) {
                return pair_energy.potential(pair.first.get(), pair.second.get());
            });
#else
        return accumulateSerial(); // fallback
#endif
    }

    double accumulateOpenMP() const
    {
        double sum = 0.0;
#pragma omp parallel for reduction(+ : sum)
        for (const auto& pair : particle_pairs) {
            sum += pair_energy.potential(pair.first.get(), pair.second.get());
        }
        return sum;
    }
};

template <RequirePairEnergy TPairEnergy>
std::unique_ptr<EnergyAccumulatorBase>
createEnergyAccumulator(const json& j, const TPairEnergy& pair_energy, double initial_value)
{
    std::unique_ptr<EnergyAccumulatorBase> accumulator;
    if (j.value("summation_policy", EnergyAccumulatorBase::Scheme::SERIAL) !=
        EnergyAccumulatorBase::Scheme::SERIAL) {
        accumulator =
            std::make_unique<DelayedEnergyAccumulator<TPairEnergy>>(pair_energy, initial_value);
        faunus_logger->debug("activated delayed energy summation");
    }
    else {
        accumulator =
            std::make_unique<InstantEnergyAccumulator<TPairEnergy>>(pair_energy, initial_value);
        faunus_logger->debug("activated instant energy summation");
    }
    accumulator->from_json(j);
    return accumulator;
}

class GroupCutoff
{
    double default_cutoff_squared = pc::max_value;
    PairMatrix<double>
        cutoff_squared; 
    Space::GeometryType& geometry; 
    friend void from_json(const json&, GroupCutoff&);
    friend void to_json(json&, const GroupCutoff&);
    void setSingleCutoff(double cutoff);

  public:
    inline bool cut(const Group& group1, const Group& group2)
    {
        if (group1.isAtomic() || group2.isAtomic()) {
            return false; // atomic groups have ill-defined mass centers
        }
        return geometry.sqdist(group1.mass_center, group2.mass_center) >=
               cutoff_squared(group1.id, group2.id);
    }

    double getCutoff(size_t id1, size_t id2) const;

    template <typename... Args> inline auto operator()(Args&&... args)
    {
        return cut(std::forward<Args>(args)...);
    }

    explicit GroupCutoff(Space::GeometryType& geometry);
};

void from_json(const json&, GroupCutoff&);
void to_json(json&, const GroupCutoff&);

template <typename TCutoff> class GroupPairingPolicy
{
  protected:
    const Space& spc; 
    TCutoff cut; 

  public:
    explicit GroupPairingPolicy(Space& spc)
        : spc(spc)
        , cut(spc.geometry)
    {
    }

    void from_json(const json& j) { Energy::from_json(j, cut); }

    void to_json(json& j) const { Energy::to_json(j, cut); }

    template <RequireEnergyAccumulator TAccumulator, typename T>
    inline void particle2particle(TAccumulator& pair_accumulator, const T& a, const T& b) const
    {
        pair_accumulator += {std::cref(a), std::cref(b)};
    }

    template <RequireEnergyAccumulator TAccumulator, typename TGroup>
    void groupInternal(TAccumulator& pair_accumulator, const TGroup& group)
    {
        const auto& moldata = group.traits();
        if (!moldata.rigid) {
            const int group_size = group.size();
            for (int i = 0; i < group_size - 1; ++i) {
                for (int j = i + 1; j < group_size; ++j) {
                    // This compound condition is faster than an outer atomic condition;
                    // tested on bulk example in GCC 9.2.
                    if (group.isAtomic() || !moldata.isPairExcluded(i, j)) {
                        particle2particle(pair_accumulator, group[i], group[j]);
                    }
                }
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename TGroup>
    void groupInternal(TAccumulator& pair_accumulator, const TGroup& group, const std::size_t index)
    {
        const auto& moldata = group.traits();
        if (!moldata.rigid) {
            if (group.isAtomic()) {
                // speed optimization: non-bonded interaction exclusions do not need to be checked
                // for atomic groups
                for (int i = 0; i < index; ++i) {
                    particle2particle(pair_accumulator, group[index], group[i]);
                }
                for (int i = index + 1; i < group.size(); ++i) {
                    particle2particle(pair_accumulator, group[index], group[i]);
                }
            }
            else {
                // molecular group
                for (int i = 0; i < index; ++i) {
                    if (!moldata.isPairExcluded(index, i)) {
                        particle2particle(pair_accumulator, group[index], group[i]);
                    }
                }
                for (int i = index + 1; i < group.size(); ++i) {
                    if (!moldata.isPairExcluded(index, i)) {
                        particle2particle(pair_accumulator, group[index], group[i]);
                    }
                }
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename TGroup, typename TIndex>
    void groupInternal(TAccumulator& pair_accumulator, const TGroup& group, const TIndex& index)
    {
        auto& moldata = group.traits();
        if (!moldata.rigid) {
            if (index.size() == 1) {
                groupInternal(pair_accumulator, group, index[0]);
            }
            else {
                // TODO investigate overhead of `index_complement` filtering;
                // TODO perhaps allow different strategies based on the index-size/group-size ratio
                auto index_complement = indexComplement(group.size(), index);
                // moved <-> static
                for (int i : index) {
                    for (int j : index_complement) {
                        if (!moldata.isPairExcluded(i, j)) {
                            particle2particle(pair_accumulator, group[i], group[j]);
                        }
                    }
                }
                // moved <-> moved
                for (auto i_it = index.begin(); i_it < index.end(); ++i_it) {
                    for (auto j_it = std::next(i_it); j_it < index.end(); ++j_it) {
                        if (!moldata.isPairExcluded(*i_it, *j_it)) {
                            particle2particle(pair_accumulator, group[*i_it], group[*j_it]);
                        }
                    }
                }
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename TGroup>
    void group2group(TAccumulator& pair_accumulator, const TGroup& group1, const TGroup& group2)
    {
        if (!cut(group1, group2)) {
            for (auto& particle1 : group1) {
                for (auto& particle2 : group2) {
                    particle2particle(pair_accumulator, particle1, particle2);
                }
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename TGroup>
    void group2group(TAccumulator& pair_accumulator, const TGroup& group1, const TGroup& group2,
                     const std::vector<std::size_t>& index1)
    {
        if (!cut(group1, group2)) {
            for (auto particle1_ndx : index1) {
                for (auto& particle2 : group2) {
                    particle2particle(pair_accumulator, *(group1.begin() + particle1_ndx),
                                      particle2);
                }
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename TGroup>
    void group2group(TAccumulator& pair_accumulator, const TGroup& group1, const TGroup& group2,
                     const std::vector<std::size_t>& index1, const std::vector<std::size_t>& index2)
    {
        if (!cut(group1, group2)) {
            if (!index2.empty()) {
                // (∁⊕group1 × ⊕group2) + (⊕group1 × ⊕group2) = group1 × ⊕group2
                group2group(pair_accumulator, group2, group1, index2);
                // + (⊕group1 × ∁⊕group2)
                auto index2_complement = indexComplement(group2.size(), index2);
                for (auto particle1_ndx : index1) {
                    for (auto particle2_ndx : index2_complement) {
                        particle2particle(pair_accumulator, group2[particle2_ndx],
                                          group1[particle1_ndx]);
                    }
                }
            }
            else if (!index1.empty()) {
                // (⊕group1 × ∁⊕group2) + (⊕group1 × ⊕group2) = ⊕group1 × group2
                group2group(pair_accumulator, group1, group2, index1);
                // + (∁⊕group1 × ⊕group2) = Ø as ⊕group2 is empty
            }
            else {
                // both indices empty hence nothing to do
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename TGroup, typename TGroups>
    void group2groups(TAccumulator& pair_accumulator, const TGroup& group, const TGroups& groups)
    {
        for (auto& other_group : groups) {
            if (&other_group != &group) {
                group2group(pair_accumulator, group, other_group);
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename TGroup, typename TGroups>
    void group2groups(TAccumulator& pair_accumulator, const TGroup& group,
                      const TGroups& group_index, const std::vector<std::size_t>& index)
    {
        for (auto other_group_ndx : group_index) {
            const auto& other_group = spc.groups[other_group_ndx];
            if (&other_group != &group) {
                group2group(pair_accumulator, group, other_group, index);
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename Tgroup>
    void group2all(TAccumulator& pair_accumulator, const Tgroup& group)
    {
        for (auto& other_group : spc.groups) {
            if (&other_group != &group) {
                group2group(pair_accumulator, group, other_group);
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename TGroup>
    void group2all(TAccumulator& pair_accumulator, const TGroup& group, const int index)
    {
        const auto& particle = group[index];
        for (auto& other_group : spc.groups) {
            if (&other_group != &group) {                      // avoid self-interaction
                if (!cut(other_group, group)) {                // check g2g cut-off
                    for (auto& other_particle : other_group) { // loop over particles in other group
                        particle2particle(pair_accumulator, particle, other_particle);
                    }
                }
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename Tgroup>
    void group2all(TAccumulator& pair_accumulator, const Tgroup& group,
                   const std::vector<std::size_t>& index)
    {
        if (index.size() == 1) {
            group2all(pair_accumulator, group, index[0]);
        }
        else {
            for (auto& other_group : spc.groups) {
                if (&other_group != &group) {
                    group2group(pair_accumulator, group, other_group, index);
                }
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename T>
    void groups2self(TAccumulator& pair_accumulator, const T& group_index)
    {
        for (auto group1_ndx_it = group_index.begin(); group1_ndx_it < group_index.end();
             ++group1_ndx_it) {
            // no such move exists that the internal energy has to be recalculated
            // groupInternal(pair_accumulator, spc.groups[*group1_ndx_it]);
            for (auto group2_ndx_it = std::next(group1_ndx_it); group2_ndx_it < group_index.end();
                 group2_ndx_it++) {
                group2group(pair_accumulator, spc.groups[*group1_ndx_it],
                            spc.groups[*group2_ndx_it]);
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename T>
    void groups2all(TAccumulator& pair_accumulator, const T& group_index)
    {
        groups2self(pair_accumulator, group_index);
        auto index_complement = indexComplement(spc.groups.size(), group_index);
        for (auto group1_ndx : group_index) {
            for (auto group2_ndx : index_complement) {
                group2group(pair_accumulator, spc.groups[group1_ndx], spc.groups[group2_ndx]);
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator> void all(TAccumulator& pair_accumulator)
    {
        for (auto group_it = spc.groups.begin(); group_it < spc.groups.end(); ++group_it) {
            groupInternal(pair_accumulator, *group_it);
            for (auto other_group_it = std::next(group_it); other_group_it < spc.groups.end();
                 other_group_it++) {
                group2group(pair_accumulator, *group_it, *other_group_it);
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator, typename TCondition>
    void all(TAccumulator& pair_accumulator, TCondition condition)
    {
        for (auto group_it = spc.groups.begin(); group_it < spc.groups.end(); ++group_it) {
            if (condition(*group_it)) {
                groupInternal(pair_accumulator, *group_it);
            }
            for (auto other_group_it = std::next(group_it); other_group_it < spc.groups.end();
                 other_group_it++) {
                group2group(pair_accumulator, *group_it, *other_group_it);
            }
        }
    }
};

template <typename TPolicy> class GroupPairing
{
    const Space& spc;
    TPolicy pairing;

  protected:
    template <RequireEnergyAccumulator TAccumulator>
    void accumulateGroup(TAccumulator& pair_accumulator, const Change& change)
    {
        const auto& change_data = change.groups.at(0);
        const auto& group = spc.groups.at(change_data.group_index);
        if (change_data.relative_atom_indices.size() == 1) {
            // faster algorithm if only a single particle moves
            pairing.group2all(pair_accumulator, group, change_data.relative_atom_indices[0]);
            if (change_data.internal) {
                pairing.groupInternal(pair_accumulator, group,
                                      change_data.relative_atom_indices[0]);
            }
        }
        else {
            const bool change_all =
                change_data.relative_atom_indices.empty(); // all particles or only their subset?
            if (change_all) {
                pairing.group2all(pair_accumulator, group);
                if (change_data.internal) {
                    pairing.groupInternal(pair_accumulator, group);
                }
            }
            else {
                pairing.group2all(pair_accumulator, group, change_data.relative_atom_indices);
                if (change_data.internal) {
                    pairing.groupInternal(pair_accumulator, group,
                                          change_data.relative_atom_indices);
                }
            }
        }
    }

    template <RequireEnergyAccumulator TAccumulator>
    void accumulateSpeciation(TAccumulator& pair_accumulator, const Change& change)
    {
        assert(change.matter_change);
        const auto& moved = change.touchedGroupIndex(); // index of moved groups
        const auto fixed = indexComplement(spc.groups.size(), moved) |
                           ranges::to<std::vector>; // index of static groups
        auto filter_active = [](int size) {
            return std::views::filter([size](const auto i) { return i < size; });
        };

        // loop over all changed groups
        for (auto change_group1_it = change.groups.begin(); change_group1_it < change.groups.end();
             ++change_group1_it) {
            const auto& group1 = spc.groups.at(change_group1_it->group_index);
            // filter only active particles
            const auto index1 = change_group1_it->relative_atom_indices |
                                filter_active(group1.size()) | ranges::to<std::vector>;
            if (!index1.empty()) {
                // particles added into the group: compute (changed group) <-> (static group)
                pairing.group2groups(pair_accumulator, group1, fixed, index1);
            }
            // loop over successor changed groups (hence avoid double counting group1×group2 and
            // group2×group1)
            for (auto change_group2_it = std::next(change_group1_it);
                 change_group2_it < change.groups.end(); ++change_group2_it) {
                const auto& group2 = spc.groups.at(change_group2_it->group_index);
                const auto index2 = change_group2_it->relative_atom_indices |
                                    filter_active(group2.size()) | ranges::to<std::vector>;
                if (!index1.empty() || !index2.empty()) {
                    // particles added into one or other group: compute (changed group) <-> (changed
                    // group)
                    pairing.group2group(pair_accumulator, group1, group2, index1, index2);
                }
            }
            if (!index1.empty() && !molecules.at(group1.id).rigid) {
                // compute internal energy in the changed group
                if (change_group1_it->all) {
                    pairing.groupInternal(pair_accumulator, group1);
                }
                else {
                    pairing.groupInternal(pair_accumulator, group1, index1);
                };
            }
        }
    }

  public:
    template <RequireEnergyAccumulator TAccumulator>
    void accumulate(TAccumulator& pair_accumulator, const Change& change)
    {
        assert(std::is_sorted(change.groups.begin(), change.groups.end()));
        if (change.everything) {
            pairing.all(pair_accumulator);
        }
        else if (change.volume_change) {
            // sum all interaction energies except the internal energies of incompressible molecules
            pairing.all(pair_accumulator, [](auto& group) {
                return group.isAtomic() || group.traits().compressible;
            });
        }
        else if (!change.matter_change) {
            if (change.groups.size() == 1) {
                // if only a single group changes use faster algorithm and optionally add the
                // internal energy
                accumulateGroup(pair_accumulator, change);
            }
            else {
                // if multiple groups move, no internal energies are computed
                const auto& moved = change.touchedGroupIndex(); // index of moved groups
                pairing.groups2all(pair_accumulator, moved);
            }
        }
        else { // change.dN
            accumulateSpeciation(pair_accumulator, change);
        }
    }

    GroupPairing(Space& spc)
        : spc(spc)
        , pairing(spc)
    {
    }

    void from_json(const json& j) { pairing.from_json(j); }

    void to_json(json& j) const { pairing.to_json(j); }

    // FIXME a temporal fix for non-refactorized NonbondedCached
    template <typename Accumulator>
    void group2group(Accumulator& pair_accumulator, const Space::GroupType& group1,
                     const Space::GroupType& group2)
    {
        pairing.group2group(std::forward<Accumulator&>(pair_accumulator),
                            std::forward<const Space::GroupType&>(group1),
                            std::forward<const Space::GroupType&>(group2));
    }
};

class NonbondedBase : public EnergyTerm
{
  public:
    virtual double particleParticleEnergy(const Particle& particle1, const Particle& particle2) = 0;
    virtual double groupGroupEnergy(const Group& group1, const Group& group2) = 0;
};

template <RequirePairEnergy TPairEnergy, typename TPairingPolicy>
class Nonbonded : public NonbondedBase
{
  protected:
    const Space& spc;        
    TPairEnergy pair_energy; 
    TPairingPolicy
        pairing; 
    std::shared_ptr<EnergyAccumulatorBase>
        energy_accumulator; 

  public:
    Nonbonded(const json& j, Space& spc, BasePointerVector<EnergyTerm>& pot)
        : spc(spc)
        , pair_energy(spc, pot)
        , pairing(spc)
    {
        name = "nonbonded";
        from_json(j);
        energy_accumulator = createEnergyAccumulator(j, pair_energy, 0.0);
        energy_accumulator->reserve(spc.numParticles()); // attempt to reduce memory fragmentation
    }

    double particleParticleEnergy(const Particle& particle1, const Particle& particle2) override
    {
        return pair_energy(particle1, particle2);
    }

    double groupGroupEnergy(const Group& group1, const Group& group2) override
    {
        InstantEnergyAccumulator<TPairEnergy> accumulator(pair_energy);
        pairing.group2group(accumulator, group1, group2);
        return static_cast<double>(accumulator);
    }

    void from_json(const json& j)
    {
        pair_energy.from_json(j);
        pairing.from_json(j);
    }

    void to_json(json& j) const override
    {
        pair_energy.to_json(j);
        pairing.to_json(j);
        energy_accumulator->to_json(j);
    }

    double energy(const Change& change) override
    {
        energy_accumulator->clear();
        // down-cast to avoid slow, virtual function calls:
        if (auto ptr = std::dynamic_pointer_cast<InstantEnergyAccumulator<TPairEnergy>>(
                energy_accumulator)) {
            pairing.accumulate(*ptr, change);
        }
        else if (auto ptr = std::dynamic_pointer_cast<DelayedEnergyAccumulator<TPairEnergy>>(
                     energy_accumulator)) {
            pairing.accumulate(*ptr, change);
        }
        else {
            pairing.accumulate(*energy_accumulator, change);
        }
        return static_cast<double>(*energy_accumulator);
    }

    void force(std::vector<Point>& forces) override
    {
        // just a temporary hack; perhaps better to allow PairForce instead of the PairEnergy
        // template
        assert(forces.size() == spc.particles.size() &&
               "the forces size must match the particle size");
        for (size_t i = 0; i < spc.particles.size() - 1; ++i) {
            for (size_t j = i + 1; j < spc.particles.size(); ++j) {
                const Point f = pair_energy.force(spc.particles[i], spc.particles[j]);
                forces[i] += f;
                forces[j] -= f;
            }
        }
    }
};

template <RequirePairEnergy TPairEnergy, typename TPairingPolicy>
class NonbondedCached : public Nonbonded<TPairEnergy, TPairingPolicy>
{
    using Base = Nonbonded<TPairEnergy, TPairingPolicy>;
    using TAccumulator = InstantEnergyAccumulator<TPairEnergy>;
    Eigen::MatrixXf energy_cache;
    using Base::spc;

    template <typename TGroup> double g2g(const TGroup& g1, const TGroup& g2)
    {
        int i = &g1 - spc.groups.data();
        int j = &g2 - spc.groups.data();
        if (j < i) {
            std::swap(i, j);
        }
        if (EnergyTerm::state ==
            EnergyTerm::MonteCarloState::TRIAL) { // if this is from the trial system
            TAccumulator energy_accumulator(Base::pair_energy);
            Base::pairing.group2group(energy_accumulator, g1, g2);
            energy_cache(i, j) = static_cast<double>(energy_accumulator); // update the cache
        }
        return energy_cache(i, j); // return (cached) value
    }

    template <typename TGroup>
    double g2g(const TGroup& g1, const TGroup& g2,
               [[maybe_unused]] const std::vector<std::size_t>& index)
    {
        // index not implemented
        return g2g(g1, g2);
    }

  public:
    NonbondedCached(const json& j, Space& spc, BasePointerVector<EnergyTerm>& pot)
        : Base(j, spc, pot)
    {
        Base::name += "EM";
        init();
    }

    void init() override
    {
        const auto groups_size = spc.groups.size();
        energy_cache.resize(groups_size, groups_size);
        energy_cache.setZero();
        TAccumulator u(Base::pair_energy);
        for (auto i = 0; i < groups_size - 1; ++i) {
            for (auto j = i + 1; j < groups_size; ++j) {
                u = 0.0;
                Base::pairing.group2group(u, spc.groups.at(i), spc.groups.at(j));
                energy_cache(i, j) = static_cast<double>(u);
            }
        }
    }

    double energy(const Change& change) override
    {
        if (!change) {
            return 0.0;
        }
        // Only g2g may be called there to compute (and cache) energy!
        double energy_sum = 0.0;
        if (change.everything || change.volume_change) {
            for (auto i = spc.groups.begin(); i < spc.groups.end(); ++i) {
                for (auto j = std::next(i); j < Base::spc.groups.end(); ++j) {
                    energy_sum += g2g(*i, *j);
                }
            }
            return energy_sum;
        };

        if (change.groups.size() == 1) { // if exactly ONE molecule is changed
            const auto& d = change.groups[0];
            const auto& g1 = spc.groups.at(d.group_index);
            for (auto g2_it = spc.groups.begin(); g2_it < spc.groups.end(); ++g2_it) {
                if (&g1 != &(*g2_it)) {
                    energy_sum += g2g(g1, *g2_it, d.relative_atom_indices);
                }
            }
            return energy_sum;
        }

        // many molecules are changed:

        auto moved = change.touchedGroupIndex(); // index of moved groups
        // moved<->moved
        if (change.moved_to_moved_interactions) {
            for (auto i = moved.begin(); i < moved.end(); ++i) {
                for (auto j = std::next(i); j < moved.end(); ++j) {
                    energy_sum += g2g(spc.groups[*i], spc.groups[*j]);
                }
            }
        }
        // moved<->static
#if true
        // classic version
        const auto fixed =
            indexComplement(spc.groups.size(), moved) | ranges::to_vector; // static groups
        for (auto i : moved) {
            for (auto j : fixed) {
                energy_sum += g2g(spc.groups[i], spc.groups[j]);
            }
        }
#else
        // OMP-ready version
        auto fixed = indexComplement(spc.groups.size(), moved) |
                     ranges::to<std::vector>; // index of static groups
        const size_t moved_size = moved.size();
        const size_t fixed_size = fixed.size();
        for (auto i = 0; i < moved_size; ++i) {
            for (auto j = 0; j < fixed_size; ++j) {
                energy_sum += g2g(spc.groups[moved[i]], spc.groups[fixed[j]]);
            }
        }
#endif
        return energy_sum;
    }

    void sync(EnergyTerm* base_ptr, const Change& change) override
    {
        auto other = dynamic_cast<decltype(this)>(base_ptr);
        assert(other);
        if (change.everything || change.volume_change) {
            energy_cache.triangularView<Eigen::StrictlyUpper>() =
                (other->energy_cache).template triangularView<Eigen::StrictlyUpper>();
        }
        else {
            for (const auto& d : change.groups) {
                for (int i = 0; i < d.group_index; i++) {
                    energy_cache(i, d.group_index) = other->energy_cache(i, d.group_index);
                }
                for (size_t i = d.group_index + 1; i < spc.groups.size(); i++) {
                    energy_cache(d.group_index, i) = other->energy_cache(d.group_index, i);
                }
            }
        }
    }
};

#ifdef ENABLE_FREESASA

class FreeSASAEnergy : public EnergyTerm
{
  private:
    std::vector<double> positions; 
    std::vector<double> radii;     
    std::vector<double> sasa;      

    const Space& spc;
    double cosolute_molarity = 0.;                       
    std::unique_ptr<freesasa_parameters_fwd> parameters; 
    Average<double> mean_surface_area;

    void to_json(json& j) const override;
    void sync(EnergyTerm* energybase_ptr, const Change& change) override;
    void updateSASA(const Change& change);
    void init() override;

    template <typename Tfirst, typename Tend>
    void updateRadii(Tfirst begin, Tend end, [[maybe_unused]] const Change& change)
    {
        const auto number_of_particles = std::distance(begin, end);
        radii.clear();
        radii.reserve(number_of_particles);
        std::transform(begin, end, std::back_inserter(radii),
                       [](const Particle& particle) { return particle.traits().sigma * 0.5; });
    }

    template <typename Tfirst, typename Tend>
    void updatePositions(Tfirst begin, Tend end, [[maybe_unused]] const Change& change)
    {
        const auto number_of_particles = std::distance(begin, end);
        positions.clear();
        positions.reserve(3 * number_of_particles);
        for (const auto& particle : spc.activeParticles()) {
            const auto* xyz = particle.pos.data();
            positions.insert(positions.end(), xyz, xyz + 3);
        }
    }

  public:
    FreeSASAEnergy(const Space& spc, double cosolute_molarity, double probe_radius);
    FreeSASAEnergy(const json& j, const Space& spc);
    double energy(const Change& change) override;

    const std::vector<double>& getAreas() const { return sasa; }
}; 
#endif

class SASAEnergyReference : public EnergyTerm
{
  private:
    void to_json(json& j) const override;
    void sync(EnergyTerm* energybase_ptr, const Change& change) override;

  protected:
    void init() override;
    using index_type = size_t;
    const Space& spc;               
    std::vector<double> areas;      
    double cosolute_molarity = 0.0; 
    std::unique_ptr<SASA::SASABase>
        sasa; 

    inline auto indexOf(const Particle& particle) const
    {
        return static_cast<index_type>(std::addressof(particle) -
                                       std::addressof(spc.particles.at(0)));
    }

  public:
    SASAEnergyReference(const Space& spc, double cosolute_molarity, double probe_radius,
                        int slices_per_atom = 25, bool dense_container = true);
    SASAEnergyReference(const json& j, const Space& spc);
    const std::vector<double>& getAreas() const;
    double energy(const Change& change) override;
};

class SASAEnergy : public SASAEnergyReference
{
  private:
    std::vector<std::vector<index_type>>
        current_neighbours; 
    std::vector<index_type>
        changed_indices; 

    void sync(EnergyTerm* energybase_ptr, const Change& change) override;
    void init() override;

    void updateChangedIndices(const Change& change);
    void insertChangedNeighboursOf(index_type index, std::set<index_type>& target_indices) const;

  public:
    SASAEnergy(const Space& spc, double cosolute_molarity, double probe_radius,
               int slices_per_atom = 25, bool dense_container = true);
    SASAEnergy(const json& j, const Space& spc);
    double energy(const Change& change) override;
};

class Example2D : public EnergyTerm
{
  private:
    bool use_2d = true;        // Set to false to apply energy only along x (as by the book)
    double scale_energy = 1.0; // effective temperature
    const Point& particle;     // reference to 1st particle in the system
    void to_json(json& j) const override;

  public:
    Example2D(const json& j, Space& spc);
    double energy(const Change& change) override;
};

class CustomGroupGroup : public EnergyTerm
{
  private:
    const Space& spc;
    MoleculeData::index_type molid1;
    MoleculeData::index_type molid2;
    std::map<MoleculeData::index_type, Average<double>> mean_charges;
    std::unique_ptr<ExprFunction<double>> expr;
    json json_input_backup; // initial json input

    struct Properties
    { // storage for particle properties
        double mean_charge1 = 0;
        double mean_charge2 = 0;
        double mass_center_separation = 0;
    };

    Properties properties;

    void to_json(json& j) const override;
    void setParameters(const Group& group1, const Group& group2);

  public:
    CustomGroupGroup(const json& j, const Space& spc);
    double energy(const Change& change) override;
};

class Hamiltonian : public EnergyTerm, public BasePointerVector<EnergyTerm>
{
  private:
    double maximum_allowed_energy = pc::infty; 
    std::vector<double>
        latest_energies;         
    decltype(vec)& energy_terms; 
    void
    addEwald(const json& j,
             Space& spc); 
    void checkBondedMolecules() const; 
    void to_json(json& j) const override;
    void force(PointVector& forces) override;
    std::unique_ptr<EnergyTerm> createEnergy(Space& spc, const std::string& name, const json& j);

  public:
    Hamiltonian(Space& spc, const json& j);
    void init() override;
    void updateState(const Change& change) override;
    void sync(EnergyTerm* other_hamiltonian, const Change& change) override;
    double energy(const Change& change) override; 
    const std::vector<double>&
    latestEnergies() const; 
};
} // namespace Energy
} // namespace Faunus