celllistimpl.h

#pragma once

#include "core.h"
#include "celllist.h"
#include <set>
#include <map>
#include <algorithm>
#include <cassert>
#include <range/v3/range/conversion.hpp>
#include <spdlog/spdlog.h>

namespace Faunus::CellList {

namespace Grid {

template <typename TGridType> struct AbstractGrid : private TGridType
{
    using GridType = TGridType;
    using typename GridType::CellCoord;
    using typename GridType::CellIndex;
    using typename GridType::Point;
    using typename GridType::SpaceAxis;

    virtual CellCoord coordinates(CellIndex index) const = 0;
    virtual CellIndex index(const CellCoord& coordinates) const = 0;
    virtual CellCoord coordinatesAt(const Point& position) const = 0;
    virtual CellIndex indexAt(const Point& position) const = 0;
    virtual bool isCell(const CellCoord& coordinates) const = 0;
    virtual bool isCellAt(const Point& point) const = 0;
    virtual CellIndex size() const = 0;
    virtual ~AbstractGrid() = default;
};

template <typename TGridType> struct AbstractNeighborGrid : public virtual AbstractGrid<TGridType>
{
    using GridType = TGridType;
    using typename TGridType::CellCoord;

    virtual bool isNeighborCell(const CellCoord& coordinates, const CellCoord& offset) const = 0;
    virtual ~AbstractNeighborGrid() = default;
};

template <typename TGridType> class GridBase : public virtual AbstractGrid<TGridType>
{
  public:
    using GridType = typename AbstractGrid<TGridType>::GridType;
    using typename GridType::CellCoord;
    using typename GridType::CellIndex;
    using typename GridType::Point;
    using typename GridType::SpaceAxis;

    CellCoord coordinates(const CellIndex index) const override
    {
        assert(index < size());
        return cell_index_to_coordinates.unaryExpr([index](const auto modulo) {
            return index % modulo;
        }) / cell_coordinates_to_index;
    }

    CellIndex index(const CellCoord& coordinates) const override
    {
        assert(isCell(coordinates));
        return (coordinates * cell_coordinates_to_index).sum();
    }

    CellCoord coordinatesAt(const Point& position) const override
    {
        assert(isCellAt(position));
        return (position / cell).floor().template cast<CellIndex>();
    }

    CellIndex indexAt(const Point& position) const override
    {
        return index(coordinatesAt(position));
    }

    bool isCell(const CellCoord& coordinates) const override
    {
        return (coordinates < this->getCellListEnd()).all() &&
               (coordinates >= CellCoord::Zero()).all();
    }

    bool isCellAt(const Point& point) const override
    {
        return (point < this->getBox()).all() && (point >= Point::Zero()).all();
    }

    CellIndex size() const override { return cell_list_end.prod(); };

    GridBase(const Point& box, SpaceAxis minimal_cell_dimension)
        : minimal_cell_dimension(minimal_cell_dimension)
    {
        updateCellGrid(box);
    }

  protected:
    const CellCoord& getCellListEnd() const { return cell_list_end; }

    const Point& getBox() const { return box; }

    const Point& getCell() const { return cell; }

  private:
    CellCoord cell_list_end; 
    SpaceAxis minimal_cell_dimension; 
    Point box;                        
    Point cell;                       

    CellCoord cell_coordinates_to_index;
    CellCoord cell_index_to_coordinates;

    void updateCellGrid(const Point& new_box)
    {
        box = new_box;
        cell_list_end = (box / box.min(minimal_cell_dimension)).floor().template cast<CellIndex>();
        cell = box / cell_list_end.template cast<SpaceAxis>();

        // (x, y, z, …) -> (1, x, x*y, …)
        cell_coordinates_to_index[0] = 1;
        cell_index_to_coordinates[0] = cell_list_end[0];
        for (int i = 1; i < TGridType::Dimension; ++i) {
            cell_coordinates_to_index[i] = cell_list_end[i - 1] * cell_coordinates_to_index[i - 1];
            cell_index_to_coordinates[i] = cell_list_end[i] * cell_index_to_coordinates[i - 1];
        }
    }
};

template <typename TGridType>
class FixedBoundaryGrid : public GridBase<TGridType>, virtual public AbstractNeighborGrid<TGridType>
{
    using Base = GridBase<TGridType>;

  public:
    using typename AbstractNeighborGrid<TGridType>::GridType;
    using typename GridType::CellCoord;

    bool isNeighborCell(const CellCoord& coordinates, const CellCoord& offset) const
    {
        return Base::isCell(coordinates + offset);
    }

    using Base::Base; // use parent constructors
};

template <typename TGridType>
class PeriodicBoundaryGrid : public GridBase<TGridType>,
                             virtual public AbstractNeighborGrid<TGridType>
{
    using Base = GridBase<TGridType>;

  public:
    using GridType = typename AbstractNeighborGrid<TGridType>::GridType;
    using typename GridType::CellCoord;
    using typename GridType::CellIndex;
    using typename GridType::Point;
    using typename GridType::SpaceAxis;

    CellCoord coordinates(const CellIndex index) const override { return Base::coordinates(index); }

    CellIndex index(const CellCoord& coordinates) const override
    {
        auto pbc_coordinates = coordinates;
        auto& boundary_coords = this->getCellListEnd();
        for (auto i = 0; i < pbc_coordinates.size(); ++i) {
            if (coordinates[i] < 0) {
                pbc_coordinates[i] +=
                    boundary_coords[i] * (std::abs(coordinates[i]) / boundary_coords[i] + 1);
            }
            else if (coordinates[i] >= boundary_coords[i]) {
                pbc_coordinates[i] -= boundary_coords[i] * (coordinates[i] / boundary_coords[i]);
            }
        }
        return Base::index(pbc_coordinates);
    }

    CellCoord coordinatesAt(const Point& position) const override
    {
        assert(isCellAt(position));
        return (position / this->getCell()).floor().template cast<CellIndex>();
    }

    CellIndex indexAt(const Point& position) const override
    {
        return this->index(this->coordinatesAt(position));
    }

    bool isCell(const CellCoord&) const override { return true; }

    bool isCellAt(const Point&) const override { return true; }

    bool isNeighborCell([[maybe_unused]] const CellCoord& coordinates,
                        const CellCoord& offset) const override
    {
        return (offset >= neighbours_cell_offset_min).all() &&
               (offset <= neighbours_cell_offset_max).all();
    }

    PeriodicBoundaryGrid(const Point& box, SpaceAxis cell_dimension)
        : Base(box, cell_dimension)
    {
        updateCellGrid();
    }

  private:
    CellCoord neighbours_cell_offset_min; 
    CellCoord neighbours_cell_offset_max; 

    void updateCellGrid()
    {
        const auto cell_list_middle = (this->getCellListEnd() - 1).template cast<double>() * 0.5;
        neighbours_cell_offset_min = -cell_list_middle.floor().template cast<CellIndex>();
        neighbours_cell_offset_max = cell_list_middle.ceil().template cast<CellIndex>();
    }
};

template <bool PBC, class TGridType>
using GridBoundary = typename std::conditional<PBC, PeriodicBoundaryGrid<TGridType>,
                                               FixedBoundaryGrid<TGridType>>::type;
template <bool PBC> using Grid3DBoundary = GridBoundary<PBC, Grid3DType>;
using Grid3DFixed = Grid3DBoundary<false>;
using Grid3DPeriodic = Grid3DBoundary<true>;
} // namespace Grid

namespace Container {
template <typename TMember, typename TIndex>
class DenseContainer : virtual public AbstractContainer<TMember, TIndex>
{
  public:
    using ContainerType = AbstractContainer<TMember, TIndex>;
    using typename ContainerType::Index;
    using typename ContainerType::Members;

    const Members& get(Index index) const override
    {
        assert(index >= 0 && index < indexEnd());
        return container[index];
    }

    Members& get(Index index) override
    {
        assert(index >= 0 && index < indexEnd());
        return container[index];
    }

    const Members& getEmpty() const override { return empty_set; }

    std::vector<Index> indices() const override
    {
        std::vector<Index> indices;
        indices.reserve(std::distance(container.begin(), container.end()));
        auto iterator = container.begin();
        while ((iterator = std::find_if(iterator, container.end(), [](const auto& members) {
                    return !members.empty();
                })) != container.end()) {
            indices.push_back(std::distance(iterator, container.begin()));
            std::advance(iterator, 1);
        }
        return indices;
    }

    explicit DenseContainer(std::size_t size)
    {
        if (size >= max_container_size) {
            faunus_logger->error("Size of the cell list container is too large! \n"
                                 "Try using a sparse version (i.e. dense: false)");
            throw std::runtime_error("cell list `container` size too large");
        }
        container.resize(size);
    }

  protected:
    Index indexEnd() const { return container.size(); }

  private:
    const std::size_t max_container_size = 1000000000; 
    const Members empty_set;        
    std::vector<Members> container; 
};

template <typename TMember, typename TIndex>
class SparseContainer : virtual public AbstractContainer<TMember, TIndex>
{
  public:
    using ContainerType = AbstractContainer<TMember, TIndex>;
    using typename ContainerType::Index;
    using typename ContainerType::Members;

    const Members& get(Index index) const override
    {
        assert(index >= 0 && index < indexEnd());
        try {
            return container.at(index);
        }
        catch (std::out_of_range& e) {
            return empty_set;
        }
    }

    Members& get(Index index) override
    {
        assert(index >= 0 && index < indexEnd());
        try {
            return container.at(index);
        }
        catch (std::out_of_range& e) {
            [[maybe_unused]] auto [iterator, flag] = container.emplace(index, Members{});
            return iterator->second;
        }
    }

    const Members& getEmpty() const override { return empty_set; }

    std::vector<Index> indices() const override
    {
        std::vector<Index> indices;
        indices.reserve(std::distance(container.begin(), container.end()));
        std::transform(container.begin(), container.end(), back_inserter(indices),
                       [](const auto& pair) { return pair.first; });
        return indices;
    }

    explicit SparseContainer(std::size_t size)
        : index_end(size)
    {
    }

  protected:
    Index indexEnd() const { return index_end; }

  private:
    Index index_end;         
    const Members empty_set; 
    std::unordered_map<Index, Members> container; 
};

template <class TContainer> class Container : public TContainer
{
  public:
    using ContainerType = ContainerTypeOf<TContainer>;
    using Member = MemberOf<ContainerType>;
    using Index = IndexOf<ContainerType>;

    void insert(const Member& member, const Index index_new)
    {
        assert(index_new < this->indexEnd());
        this->get(index_new).push_back(member);
    }

    void erase(const Member& member, const Index index_old)
    {
        assert(index_old < this->indexEnd());
        auto& members = this->get(index_old);
        // std::erase_if available in C++20
        auto remove_it = std::remove(members.begin(), members.end(), member);
        members.erase(remove_it, members.end());
        // end of std::erase_if
    }

    void move(const Member& member, const Index index_old, const Index index_new)
    {
        erase(member, index_old);
        insert(member, index_new);
    }

    using TContainer::TContainer; // use parent constructor
};

} // namespace Container

template <class TContainer, class TGrid>
class CellListBase
    : protected TGrid,
      protected TContainer,
      virtual public AbstractImmutableCellList<ContainerTypeOf<TContainer>, GridTypeOf<TGrid>>
{
  public:
    using AbstractCellList =
        AbstractImmutableCellList<ContainerTypeOf<TContainer>, GridTypeOf<TGrid>>;
    using Grid = TGrid;
    using Container = TContainer;
    using typename AbstractCellList::CellCoord; 
    using typename AbstractCellList::CellIndex; 
    using typename AbstractCellList::GridType;  
    using typename AbstractCellList::Member;    
    using typename AbstractCellList::Members;   

    const Members& getMembers(const CellCoord& cell_coordinates) override
    {
        return this->get(this->index(cell_coordinates));
    }

    const Members& getNeighborMembers(const CellCoord& cell_coordinates,
                                      const CellCoord& cell_offset) override
    {
        if (this->isNeighborCell(cell_coordinates, cell_offset)) {
            return getMembers(cell_coordinates + cell_offset);
        }
        else {
            return getEmpty();
        }
    }

    std::vector<CellCoord> getCells() const override
    {
        const auto indices = this->indices();
        return std::views::all(indices) |
               std::views::transform([this](auto index) { return this->coordinates(index); }) |
               ranges::to_vector;
    }

    CellIndex gridSize() const override { return Grid::size(); }

    const TGrid& getGrid() { return *this; }

    explicit CellListBase(const TGrid& cell_grid)
        : Grid(cell_grid)
        , Container(cell_grid.size())
    {
    }

    CellListBase(const typename TGrid::Point& box, const typename TGrid::SpaceAxis cell_dimension)
        : Grid(box, cell_dimension)
        , Container(Grid::size())
    {
    }

  protected:
    //    const auto& getMembers(const CellIndex cell_index) const { return this->get(cell_index); }
    //    auto& getMembers(const CellIndex cell_index) { return this->get(cell_index); }
    //    const auto& getEmptyCell() const { return this->getEmpty(); }

    using Container::erase;
    using Container::get;
    using Container::getEmpty;
    using Container::insert;
    using Container::move;
};

template <typename TMember, class TGrid,
          template <typename, typename> class TCellList = CellListBase,
          template <typename, typename> class TContainer = Container::DenseContainer>
using CellListType =
    TCellList<Container::Container<TContainer<TMember, Grid::IndexOf<TGrid>>>, TGrid>;

template <class TContainer, class TGrid>
class SortableCellList
    : public CellListBase<TContainer, TGrid>,
      virtual public AbstractSortableCellList<ContainerTypeOf<TContainer>, GridTypeOf<TGrid>>
{
    using TBase = CellListBase<TContainer, TGrid>;

  public:
    using AbstractCellList = typename AbstractSortableCellList<ContainerTypeOf<TContainer>,
                                                               GridTypeOf<TGrid>>::AbstractCellList;
    using Grid = TGrid;
    using Container = TContainer;
    using typename AbstractCellList::CellCoord; 
    using typename AbstractCellList::CellIndex; 
    using typename AbstractCellList::Member;    
    using typename AbstractCellList::Members;   

    const Members& getSortedMembers(const CellCoord& cell_coordinates) override
    {
        return getSorted(this->index(cell_coordinates));
    }

    const Members& getNeighborSortedMembers(const CellCoord& cell_coordinates,
                                            const CellCoord& cell_offset) override
    {
        if (this->isNeighborCell(cell_coordinates, cell_offset)) {
            return getSortedMembers(cell_coordinates + cell_offset);
        }
        else {
            return this->getEmpty();
        }
    }

    explicit SortableCellList(const GridOf<TBase>& cell_grid)
        : TBase(cell_grid)
    {
        is_sorted.resize(this->size(), false);
    }

    SortableCellList(const typename GridOf<TBase>::Point& box,
                     const typename GridOf<TBase>::SpaceAxis cell_dimension)
        : TBase(box, cell_dimension)
    {
        is_sorted.resize(this->size(), false);
    }

  protected:
    void insert(const Member& member, const CellIndex& new_cell_index)
    {
        TBase::insert(member, new_cell_index);
        if (is_sorted[new_cell_index]) {
            is_sorted[new_cell_index] = false;
        }
    }

    void remove(const Member& member, const CellIndex& old_cell_index)
    {
        TBase::remove(member, old_cell_index);
        if (is_sorted[old_cell_index]) {
            is_sorted[old_cell_index] = false;
        }
    }

    const Members& getSorted(const CellIndex cell_index)
    {
        auto& members = this->get(cell_index);
        if (!is_sorted[cell_index]) {
            std::sort(members.begin(), members.end());
            is_sorted[cell_index] = true;
        }
        assert(std::is_sorted(members.begin(), members.end()));
        return members;
    }

  private:
    // TODO perhaps a map for sparse containers
    std::vector<bool> is_sorted;
};

template <class TBase> class CellListReverseMap : public TBase
{
  public:
    using typename TBase::AbstractCellList;     
    using typename AbstractCellList::CellCoord; 
    using typename AbstractCellList::CellIndex; 
    using typename AbstractCellList::Member;    
    using typename TBase::Container;
    using typename TBase::Grid;

    void addMember(const Member& member, const CellCoord& new_cell_coordinates)
    {
        this->insert(member, this->index(new_cell_coordinates));
    }

    void removeMember(const Member& member)
    {
        const auto old_cell_index = member2cell.at(member);
        this->erase(member, old_cell_index);
        member2cell.erase(member);
    }

    void updateMember(const Member& member, const CellCoord& new_cell_coordinates)
    {
        this->update(member, this->index(new_cell_coordinates));
    }

    bool containsMember(const Member& member) { return member2cell.count(member) != 0; }

    template <typename T> void importMembers(CellListReverseMap& source, const T& members)
    {
        for (const auto& member : members) {
            insert(member, source.member2cell[member]);
        }
    }

    using TBase::TBase;

  protected:
    void insert(const Member& member, const CellIndex& new_cell_index)
    {
        assert(member2cell.count(member) == 0); // FIXME C++20 contains
        TBase::insert(member, new_cell_index);
        member2cell.insert({member, new_cell_index});
    }

    void update(const Member& member, const CellIndex& new_cell_index)
    {
        const auto old_cell_index = member2cell.at(member);
        if (new_cell_index != old_cell_index) {
            TBase::move(member, old_cell_index, new_cell_index);
            member2cell.at(member) = new_cell_index;
        }
    }

  private:
    // TODO as a vector because size of Particle vector is constant
    std::map<Member, CellIndex> member2cell; 
};

template <class TBase> class CellListSpatial : public CellListReverseMap<TBase>
{
  public:
    using typename CellListReverseMap<TBase>::AbstractCellList; 
    using typename CellListReverseMap<TBase>::Grid;
    using typename CellListReverseMap<TBase>::Container;
    using typename AbstractCellList::Member;          
    using typename AbstractCellList::GridType::Point; 

    void insertMember(const Member& member, const Point& new_position)
    {
        const auto new_cell_index = this->indexAt(new_position);
        this->insert(member, new_cell_index);
    }

    void updateMemberAt(const Member& member, const Point& new_position)
    {
        const auto new_cell_index = this->indexAt(new_position);
        this->update(member, new_cell_index);
    }

    using CellListReverseMap<TBase>::CellListReverseMap;
};

template <typename TCellListMinuend, typename TCellListSubtrahend>
class CellListDifference
    : virtual public AbstractSortableCellList<ContainerTypeOf<ContainerOf<TCellListMinuend>>,
                                              GridTypeOf<GridOf<TCellListMinuend>>>
{
    using CellListMinuend = TCellListMinuend;
    using CellListSubtrahend = TCellListSubtrahend;

  public:
    using AbstractCellList = typename CellListMinuend::AbstractCellList;
    using Grid = typename CellListMinuend::Grid;
    using Container = typename CellListMinuend::Container;
    using typename AbstractCellList::CellCoord; 
    using typename AbstractCellList::CellIndex; 
    using typename AbstractCellList::Member;    
    using typename AbstractCellList::Members;   

    const Members& getSortedMembers(const CellCoord& cell_coordinates) override
    {
        const auto subtrahend_cell = subtrahend->getSortedMembers(cell_coordinates);
        if (std::size(subtrahend_cell) == 0) {
            // nothing to subtract
            return minuend->getSortedMembers(cell_coordinates);
        }
        else {
            const auto cell_index = minuend->getGrid().index(cell_coordinates);
            auto difference_cell_it = difference_cache.find(cell_index);
            if (difference_cell_it == difference_cache.end()) {
                // not cached yet
                auto minuend_cell = minuend->getSortedMembers(cell_coordinates);
                std::tie(difference_cell_it, std::ignore) =
                    difference_cache.emplace(cell_index, Members{});
                std::set_difference(minuend_cell.begin(), minuend_cell.end(),
                                    subtrahend_cell.begin(), subtrahend_cell.end(),
                                    std::back_inserter(difference_cell_it->second));
            }
            return difference_cell_it->second;
        }
    }

    const Members& getMembers(const CellCoord& cell_coordinates) override
    {
        if (subtrahend->getMembers(cell_coordinates).empty()) {
            // avoid difference and soring when not necessary
            return minuend->getMembers(cell_coordinates);
        }
        else {
            return getSortedMembers(cell_coordinates);
        }
    }

    const Members& getNeighborSortedMembers(const CellCoord& cell_coordinates,
                                            const CellCoord& cell_offset) override
    {
        if (minuend->getGrid().isNeighborCell(cell_coordinates, cell_offset)) {
            // avoid difference and soring when not necessary
            return getSortedMembers(cell_coordinates + cell_offset);
        }
        else {
            return minuend->getNeighborMembers(cell_coordinates, cell_offset);
        }
    }

    const Members& getNeighborMembers(const CellCoord& cell_coordinates,
                                      const CellCoord& cell_offset) override
    {
        if (minuend->getGrid().isNeighborCell(cell_coordinates, cell_offset)) {
            return getMembers(cell_coordinates + cell_offset);
        }
        else {
            return minuend->getNeighborMembers(cell_coordinates, cell_offset);
        }
    }

    CellIndex gridSize() const override { return minuend->gridSize(); }

    auto& getGrid() { return minuend->getGrid(); }

    std::vector<CellCoord> getCells() const override { return minuend->getCells(); }

    CellListDifference(std::shared_ptr<TCellListMinuend> minuend,
                       std::shared_ptr<TCellListSubtrahend> subtrahend)
        : minuend(minuend)
        , subtrahend(subtrahend)
    {
    }

  private:
    std::shared_ptr<CellListMinuend> minuend;
    std::shared_ptr<CellListSubtrahend> subtrahend;
    std::map<CellIndex, Members> difference_cache;
};

} // namespace Faunus::CellList