sasa.h

#pragma once

#ifndef FAUNUS_SASA_H
#define FAUNUS_SASA_H

#include "celllistimpl.h"
#include "particle.h"
#include <range/v3/numeric.hpp>
#include <numbers>

namespace Faunus {

class Space;

namespace Geometry {
class Chameleon;
}

struct Change;

namespace SASA {

using index_type = size_t;
using GeometryType = Geometry::Chameleon;

class SASABase
{
  public:
    struct Neighbours
    {
        std::vector<index_type> indices; 
        PointVector points;              
        index_type index;                
    };

    bool needs_syncing = false;

  protected:
    double probe_radius = 1.4;      
    std::vector<double> areas;      
    std::vector<double> sasa_radii; 
    int slices_per_atom = 20;       
    const double two_pi = 2.0 * std::numbers::pi;
    const Particle* first_particle; 

    [[nodiscard]] inline index_type indexOf(const Particle& particle) const
    {
        return static_cast<index_type>(std::addressof(particle) - first_particle);
    }

    [[nodiscard]] double calcSASAOfParticle(const Neighbours& neighbour) const;
    double exposedArcLength(std::vector<std::pair<double, double>>& arcs) const;

  public:
    [[nodiscard]] double calcSASAOfParticle(const Space& spc, const Particle& particle) const;

    template <typename TBegin, typename TEnd>
    double calcSASA(const Space& spc, TBegin begin, TEnd end) const
    {
        return ranges::accumulate(begin, end, 0.0, [&spc, this](auto& area, const auto& species) {
            return area + this->calcSASAOf(spc, species);
        });
    }

    template <typename TSpecies> double calcSASAOf(const Space& spc, const TSpecies& species) const;

    void updateSASA(const std::vector<SASABase::Neighbours>& neighbours_data,
                    const std::vector<index_type>& target_indices);

    virtual void init(const Space& spc) = 0;
    [[nodiscard]] virtual std::vector<SASABase::Neighbours>
    calcNeighbourData(const Space& spc, const std::vector<index_type>& target_indices) const = 0;
    [[nodiscard]] virtual SASABase::Neighbours
    calcNeighbourDataOfParticle(const Space& spc, index_type target_index) const = 0;
    virtual void update(const Space& spc, const Change& change) = 0;
    [[nodiscard]] const std::vector<double>& getAreas() const;
    SASABase(const Space& spc, double probe_radius, int slices_per_atom);
    virtual ~SASABase() = default;
};

class SASA : public SASABase
{
  public:
    void init(const Space& spc) override;
    [[nodiscard]] std::vector<SASABase::Neighbours>
    calcNeighbourData(const Space& spc,
                      const std::vector<index_type>& target_indices) const override;
    [[nodiscard]] SASA::Neighbours
    calcNeighbourDataOfParticle(const Space& spc, index_type target_index) const override;
    void update([[maybe_unused]] const Space& spc, [[maybe_unused]] const Change& change) override;
    SASA(const Space& spc, double probe_radius, int slices_per_atom);
    SASA(const json& j, const Space& spc);
};

template <typename CellList> class SASACellList : public SASABase
{
  private:
    using CellCoord = typename CellList::Grid::CellCoord;
    std::unique_ptr<CellList> cell_list; 
    double cell_length;                  
    std::vector<CellCoord>
        cell_offsets; 

  public:
    SASACellList(const Space& spc, double probe_radius, int slices_per_atom);
    SASACellList(const json& j, const Space& spc);
    ~SASACellList() override = default;
    void init(const Space& spc) override;
    [[nodiscard]] SASABase::Neighbours
    calcNeighbourDataOfParticle(const Space& spc, index_type target_index) const override;
    [[nodiscard]] std::vector<SASABase::Neighbours>
    calcNeighbourData(const Space& spc,
                      const std::vector<index_type>& target_indices) const override;
    void update(const Space& spc, const Change& change) override;

  private:
    template <typename TBegin, typename TEnd>
    void createCellList(TBegin begin, TEnd end, const GeometryType& geometry);
    void updateMatterChange(const Space& spc, const Change& change);
    void updatePositionsChange(const Space& spc, const Change& change);
};

using PeriodicGrid = CellList::Grid::Grid3DPeriodic;
using FixedGrid = CellList::Grid::Grid3DFixed;

template <typename TMember, typename TIndex>
using DenseContainer = CellList::Container::DenseContainer<TMember, TIndex>;
template <typename TMember, typename TIndex>
using SparseContainer = CellList::Container::SparseContainer<TMember, TIndex>;

template <class TGrid, template <typename, typename> class TContainer = DenseContainer>
using CellListType = CellList::CellListSpatial<
    CellList::CellListType<index_type, TGrid, CellList::CellListBase, TContainer>>;

using DensePeriodicCellList = CellListType<PeriodicGrid, DenseContainer>;
using DenseFixedCellList = CellListType<FixedGrid, DenseContainer>;
using SparsePeriodicCellList = CellListType<PeriodicGrid, SparseContainer>;
using SparseFixedCellList = CellListType<FixedGrid, SparseContainer>;

extern template class SASACellList<DensePeriodicCellList>;
extern template class SASACellList<DenseFixedCellList>;
extern template class SASACellList<SparsePeriodicCellList>;
extern template class SASACellList<SparseFixedCellList>;

} // namespace SASA
} // namespace Faunus

#endif // FAUNUS_SASA_H