clustermove.h

#pragma once

#include "move.h"
#include "aux/pairmatrix.h"
#include <map>

namespace Faunus {
namespace move {

class ClusterShapeAnalysis
{
  private:
    unsigned long int number_of_samples = 0;
    bool save_pqr_files = false; 
    bool shape_anisotropy_use_com =
        true; 
    std::unique_ptr<std::ostream> stream;       
    std::map<size_t, size_t> size_distribution; 
    std::map<size_t, AverageObj<Geometry::ShapeDescriptors>> shape_distribution; 
    std::map<size_t, std::ofstream> pqr_distribution; 
    decltype(pqr_distribution)::iterator
    findPQRstream(size_t cluster_size); 

    template <RequireGroups Range>
    Tensor gyrationFromMassCenterPositions(
        const Range& groups, const Point& mass_center_of_groups,
        const Geometry::BoundaryFunction boundary = [](auto&) {})
    {
        using std::views::transform;
        auto positions = groups | transform(&Group::mass_center);
        auto masses = groups | transform(&Group::mass);
        return Geometry::gyration(positions.begin(), positions.end(), masses.begin(),
                                  mass_center_of_groups, boundary);
    }

    template <RequireGroups Range>
    Tensor gyrationFromParticlePositions(
        const Range& groups, const Point& mass_center_of_groups,
        const Geometry::BoundaryFunction boundary = [](auto&) {})
    {
        using namespace ranges::cpp20::views;
        auto positions = groups | join | transform(&Particle::pos);
        auto masses = groups | join | transform(&Particle::traits) | transform(&AtomData::mw);
        return Geometry::gyration(positions.begin(), positions.end(), masses.begin(),
                                  mass_center_of_groups, boundary);
    }

    friend void to_json(json&, const ClusterShapeAnalysis&);

  public:
    template <RequireGroups Range>
    void sample(const Range& groups, const Point& mass_center_of_groups, const Space& spc)
    {
        const auto cluster_size = std::distance(groups.begin(), groups.end());
        Tensor gyration_tensor;
        if (shape_anisotropy_use_com) {
            gyration_tensor = gyrationFromMassCenterPositions(groups, mass_center_of_groups,
                                                              spc.geometry.getBoundaryFunc());
        }
        else {
            gyration_tensor = gyrationFromParticlePositions(groups, mass_center_of_groups,
                                                            spc.geometry.getBoundaryFunc());
        }
        const auto shape = Geometry::ShapeDescriptors(gyration_tensor);
        size_distribution[cluster_size]++;
        shape_distribution[cluster_size] += shape;
        if (stream) {
            const auto seed_index =
                &(*groups.begin()) - &spc.groups.at(0); // index of first molecule in cluster
            *stream << std::format("{} {} {:.3f}\n", cluster_size, seed_index,
                                   shape.relative_shape_anisotropy);
        }
        if (save_pqr_files) {
            auto it = findPQRstream(cluster_size);
            it->second << std::format("REMARK   0 Sample number = {}\n", number_of_samples)
                       << std::format("REMARK   0 Relative shape anisotropy = {:.3f}\n",
                                      shape.relative_shape_anisotropy);
            PQRWriter pqr;
            pqr.style = PQRWriter::Style::PQR;
            pqr.save(it->second, groups, spc.geometry.getLength());
        }
        ++number_of_samples;
    }

    ClusterShapeAnalysis(bool shape_anisotropy_use_com, const std::string& filename,
                         bool dump_pqr_files);
    ClusterShapeAnalysis(const json& j);
};

void to_json(json& j, const ClusterShapeAnalysis& shape);

class FindCluster
{
  private:
    const Space& spc;
    bool use_mass_center_threshold =
        true;                  
    bool single_layer = false; 
    std::vector<std::string> molecule_names; 
    std::vector<int> molids; 
    std::set<int>
        satellites; 
    PairMatrix<double, true> thresholds_squared; 

    void parseThresholds(const json& j); 
    double clusterProbability(const Group& group1, const Group& group2) const;
    void registerSatellites(const std::vector<std::string>&); 
    void updateMoleculeIndex();                               
    friend void to_json(json& j, const FindCluster& cluster);

  public:
    std::vector<size_t> molecule_index; 
    std::optional<size_t> findSeed(Random& random); 
    std::pair<std::vector<size_t>, bool> findCluster(size_t seed_index); 
    FindCluster(const Space& spc, const json& j);
};

void to_json(json& j, const FindCluster& cluster); 

class GroupMover
{
  public:
    Average<double> mean_square_displacement; 
    const double displacement_factor;         
    const Point direction;                    
    GroupMover(double displacement_factor, const Point& direction);
};

class GroupTranslator : public GroupMover
{
  public:
    Point current_displacement = Point::Zero(); 
    GroupTranslator(double displacement_factor, const Point& direction = Point::Ones());
    std::function<void(Group&)> getLambda(Geometry::BoundaryFunction boundary, Random& slump);
};

class GroupRotator : public GroupMover
{
  private:
    Eigen::Quaterniond setRandomRotation(Random& random);

  public:
    double current_displacement = 0.0; 
    GroupRotator(double displacement_factor, const Point& rotation_axis = Point::Zero());
    std::function<void(Group&)> getLambda(Geometry::BoundaryFunction boundary,
                                          const Point& rotation_origin,
                                          Random& random); 
};

class Cluster : public Move
{
  private:
    std::unique_ptr<FindCluster> find_cluster;
    std::unique_ptr<ClusterShapeAnalysis> shape_analysis;
    std::unique_ptr<GroupTranslator> translate; 
    std::unique_ptr<GroupRotator> rotate;       

    std::function<Group&(size_t)> index_to_group;
    Average<double> average_cluster_size;
    double bias_energy = 0;         
    unsigned int bias_rejected = 0; 
    unsigned int shape_analysis_interval = 0; 

    Point
    clusterMassCenter(const std::vector<size_t>& indices) const; 
    void _move(Change& change) override;                         
    double
    bias(Change& change, double old_energy,
         double new_energy) override; 
    void _reject(Change& change) override;
    void _accept(Change& change) override;
    void _to_json(json& j) const override;
    void _from_json(const json& j) override;
    void setChange(Change& change, const std::vector<size_t>& group_indices) const;
    void calculateBias(size_t seed_index, const std::vector<size_t>& cluster_index);
    void clearForMove(); 

  protected:
    Cluster(Space& spc, std::string_view name, std::string_view cite);

  public:
    explicit Cluster(Space& spc);
};

} // namespace move
} // namespace Faunus