move.h

#pragma once

#include "average.h"
#include "mpicontroller.h"
#include "space.h"
#include "io.h"
#include "smart_montecarlo.h"
#include "aux/timers.h"
#include "aux/sparsehistogram.h"
#include <range/v3/view/filter.hpp>
#include <range/v3/view/iota.hpp>
#include <range/v3/view/transform.hpp>
#include <range/v3/view/indirect.hpp>

namespace Faunus {

namespace Speciation {
class GroupDeActivator;
}

namespace Energy {
class Hamiltonian;
}

namespace move {

class MoveCollection;

class Move
{
  private:
    virtual void _move(Change& change) = 0;               
    virtual void _accept(Change& change);                 
    virtual void _reject(Change& change);                 
    virtual void _to_json(json& j) const = 0;             
    virtual void _from_json(const json& j) = 0;           
    TimeRelativeOfTotal<std::chrono::microseconds> timer; 
    TimeRelativeOfTotal<std::chrono::microseconds> timer_move; 

    friend MoveCollection;
    unsigned long number_of_accepted_moves = 0;
    unsigned long number_of_rejected_moves = 0;
    unsigned int sweep_interval = 1; 
    const std::string cite;          

  protected:
    const std::string name; 
    Space& spc;             
    int repeat = 1;
    unsigned long number_of_attempted_moves = 0; 

  public:
    static Random slump; 

    void from_json(const json& j);
    void to_json(json& j) const; 
    void move(Change& change);   
    void accept(Change& change);
    void reject(Change& change);
    void setRepeat(int repeat);
    virtual double bias(Change& change, double old_energy,
                        double new_energy); 
    Move(Space& spc, std::string_view name, std::string_view cite);
    inline virtual ~Move() = default;
    [[nodiscard]] bool isStochastic() const; 
    [[nodiscard]] const std::string& getName() const;
};

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

class ReplayMove : public Move
{
    std::unique_ptr<XTCReader> reader = nullptr; 
    TrajectoryFrame frame;                       
    bool end_of_trajectory = false;              
    // FIXME resolve always accept / always reject on the Faunus level
    const double force_accept =
        -1e12; 

    void _move(Change& change) override;
    void _to_json(json&) const override;
    void _from_json(const json&) override;
    double bias(Change&, double, double) override;

  protected:
    using Move::spc;
    ReplayMove(Space& spc, const std::string& name, const std::string& cite);

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

class [[maybe_unused]] AtomicSwapCharge : public Move
{
    int molid = -1;
    double pKa, pH;
    Average<double> msqd; // mean squared displacement
    double _sqd, _bias;   // squared displament
    std::string molname;  // name of molecule to operate on
    Change::GroupChange cdata;

    void _to_json(json&) const override;
    void _from_json(const json&) override; 
    typename ParticleVector::iterator randomAtom();
    void _move(Change& change) override;
    double bias(Change&, double,
                double) override; 
    void _accept(Change&) override;
    void _reject(Change&) override;

  protected:
    using Move::spc;
    AtomicSwapCharge(Space& spc, const std::string& name, const std::string& cite);

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

class AtomicTranslateRotate : public Move
{
    ParticleVector::const_iterator latest_particle; 
    const Energy::Hamiltonian& hamiltonian;         
    std::map<int, SparseHistogram<>>
        energy_histogram;               
    double energy_resolution = 0.0;     
    double latest_displacement_squared; 
    double default_dp = 0.0;            
    double default_dprot = 0.0;         
    void sampleEnergyHistogram();       
    void saveHistograms();              
    void checkMassCenter(
        Space::GroupType& group) const; 
    void
    groupToDisk(const Space::GroupType& group) const; 

  protected:
    int molid = -1;                           
    Point directions = {1, 1, 1};             
    Average<double> mean_square_displacement; 
    std::string molecule_name;                
    Change::GroupChange cdata;                

    void _to_json(json&) const override;
    void _from_json(const json&) override; 
    ParticleVector::iterator randomAtom(); 

    virtual void translateParticle(ParticleVector::iterator, double); 
    void _move(Change&) override;
    void _accept(Change&) override;
    void _reject(Change&) override;

    AtomicTranslateRotate(Space& spc, const Energy::Hamiltonian& hamiltonian,
                          const std::string& name, const std::string& cite);

  public:
    AtomicTranslateRotate(Space& spc, const Energy::Hamiltonian& hamiltonian);
    ~AtomicTranslateRotate() override;
};

/*
template<typename Tspace>
   class Atomic2dTranslateRotate : public AtomicTranslateRotate {
       protected:
           typedef AtomicTranslateRotate base;
           using base::spc;

           void translateParticle(Space::Tpvec::iterator p, double dp) override {
               auto &g = spc.groups[base::cdata.index];
               Point oldpos = p->pos;

               Point rtp = xyz2rtp(p->pos); // Get the spherical coordinates of the particle
               double slump_theta = dp * (base::slump() - 0.5); // Get random theta-move
               double slump_phi = dp * (base::slump() - 0.5);   // Get random phi-move

               double scalefactor_theta = spc.geo.getRadius() * sin(rtp.z()); // Scale-factor for
theta double scalefactor_phi = spc.geo.getRadius();                  // Scale-factor for phi

               Point theta_dir = Point(-sin(rtp.y()), cos(rtp.y()), 0); // Unit-vector in
theta-direction Point phi_dir = Point(cos(rtp.y()) * cos(rtp.z()), sin(rtp.y()) * cos(rtp.z()),
                                     -sin(rtp.z())); // Unit-vector in phi-direction
               Point xyz = oldpos + scalefactor_theta * theta_dir * slump_theta +
                           scalefactor_phi * phi_dir * slump_phi; // New position
               p->pos = spc.geo.getRadius() * xyz / xyz.norm();   // Convert to cartesian
coordinates

               spc.geo.boundary(p->pos);
               base::_sqd = spc.geo.sqdist(oldpos, p->pos); // squared displacement
               if (not g.atomic) {                          // recalc mass-center for non-molecular
groups g.cm = Geometry::massCenter(g.begin(), g.end(), spc.geo.getBoundaryFunc(), -g.cm);

#ifndef NDEBUG
                   Point cmbak = g.cm;                             // backup mass center
                   g.translate(-cmbak, spc.geo.getBoundaryFunc()); // translate to {0,0,0}
                   double should_be_zero = spc.geo.sqdist({0, 0, 0}, Geometry::massCenter(g.begin(),
g.end())); if (should_be_zero > 1e-6) throw std::runtime_error("atomic move too large"); else
g.translate(cmbak, spc.geo.getBoundaryFunc()); #endif
               }
           }

       public:
           Atomic2dTranslateRotate(Tspace &spc, std::string name, std::string cite) : base(spc,
name, cite) {} explicit Atomic2dTranslateRotate(Tspace &spc) : Atomic2dTranslateRotate(spc,
"transrot 2d", "") {}
   };*/

class TranslateRotate : public Move
{
  protected:
    using OptionalGroup = std::optional<std::reference_wrapper<Space::GroupType>>;
    int molid = -1; 
    void _to_json(json& j) const override;
    TranslateRotate(Space& spc, std::string name, std::string cite);

  private:
    Average<double> mean_squared_displacement;
    Average<double> mean_squared_rotation_angle;

    double latest_displacement_squared = 0.0;
    double latest_rotation_angle_squared = 0.0;
    double translational_displacement = 0.0;   
    double rotational_displacement = 0.0;      
    Point translational_direction = {1, 1, 1}; 
    Point fixed_rotation_axis = {0, 0, 0};     

    virtual OptionalGroup findRandomMolecule();
    double translateMolecule(Space::GroupType& group);
    double rotateMolecule(Space::GroupType& group);
    void checkMassCenter(const Space::GroupType& group) const; // sanity check of move

    void _from_json(const json& j) override;
    void _move(Change& change) override;
    void _accept(Change&) override;
    void _reject(Change&) override;

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

class SmarterTranslateRotate : public TranslateRotate
{
  private:
    SmarterMonteCarlo::MoveSupport<Group> smartmc;
    void _to_json(json& j) const override;
    OptionalGroup findRandomMolecule() override;
    double bias(Change& change, double old_energy, double new_energy) override;

  public:
    SmarterTranslateRotate(Space& spc, const json& j);
};

class ConformationSwap : public Move
{
  public:
    enum class CopyPolicy
    {
        ALL,
        POSITIONS,
        CHARGES,
        PATCHES,
        INVALID
    }; 
  private:
    CopyPolicy copy_policy;
    RandomInserter inserter;
    int molid = -1; 
    void copyConformation(ParticleVector& source_particle,
                          ParticleVector::iterator destination) const;
    void _to_json(json& j) const override;
    void _from_json(const json& j) override;
    void _move(Change& change) override;
    void setRepeat();                   
    void checkConformationSize() const; 
    void checkMassCenterDrift(const Point& old_mass_center,
                              const ParticleVector& particles); 
    void registerChanges(Change& change,
                         const Space::GroupType& group) const; 
    ConformationSwap(Space& spc, const std::string& name, const std::string& cite);

  public:
    explicit ConformationSwap(Space& spc);
}; // end of conformation swap move

NLOHMANN_JSON_SERIALIZE_ENUM(ConformationSwap::CopyPolicy,
                             {{ConformationSwap::CopyPolicy::INVALID, nullptr},
                              {ConformationSwap::CopyPolicy::ALL, "all"},
                              {ConformationSwap::CopyPolicy::PATCHES, "patches"},
                              {ConformationSwap::CopyPolicy::POSITIONS, "positions"},
                              {ConformationSwap::CopyPolicy::CHARGES, "charges"}})

class VolumeMove : public Move
{
  protected:
    Geometry::VolumeMethod volume_scaling_method = Geometry::VolumeMethod::ISOTROPIC;
    double logarithmic_volume_displacement_factor = 0.0;
    double old_volume = 0.0;
    double new_volume = 0.0;
    void _move(Change& change) override;
    void _from_json(const json& j) override;

  private:
    Average<double> mean_volume;
    Average<double> mean_square_volume_change;

    virtual void setNewVolume();
    void _to_json(json& j) const override;
    void _accept(Change& change) override;
    void _reject(Change& change) override;

  public:
    VolumeMove(Space& spc, std::string_view name);
    explicit VolumeMove(Space& spc);
}; // end of VolumeMove

class ChargeMove : public Move
{
  private:
    Average<double> mean_squared_charge_displacement;
    Change::GroupChange group_change;

    void _move(Change& change) override;
    void _accept(Change&) override;
    void _reject(Change&) override;
    void _from_json(const json& j) override;
    [[nodiscard]] virtual double getChargeDisplacement(const Particle& particle) const;
    ChargeMove(Space& spc, std::string_view name, std::string_view cite);

  protected:
    double max_charge_displacement = 0.0;
    double charge_displacement = 0.0;
    ParticleVector::size_type particle_index;
    void _to_json(json& j) const override;

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

class QuadraticChargeMove : public ChargeMove
{
  private:
    Average<double> mean_bias;
    void _to_json(json& j) const override;
    [[nodiscard]] double getChargeDisplacement(const Particle& particle) const override;
    double bias(Change& change, double old_energy, double new_energy) override;

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

class ChargeTransfer : public Move
{
  private:
    Average<double> msqd; // mean squared displacement
    double dq = 0, deltaq = 0;

    struct moldata
    {
        double charges = 0;
        double moves = 0;
        size_t numOfAtoms = 0;
        int id = 0;
        std::string molname;
        std::vector<double> min, max;
        std::vector<double> molrange;
        std::vector<double> ratio;
        std::vector<double> changeQ;
        Change::GroupChange cdata;
    };

    moldata mol1, mol2;

    double sumTemp = 0;
    int i = 0;

    void _to_json(json& j) const override;
    void _from_json(const json& j) override;
    void _move(Change& change) override;
    void _accept(Change&) override;
    void _reject(Change&) override;

  protected:
    using Move::spc;
    ChargeTransfer(Space& spc, const std::string& name, const std::string& cite);

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

class QuadrantJump : public Move
{
  private:
    int molid = -1;
    Point dir = {1, 1, 1};
    std::vector<size_t> index;
    double _sqd;          // squared displacement
    Average<double> msqd; // mean squared displacement

    void _to_json(json& j) const override;
    void _from_json(const json& j) override; 
    void _move(Change& change) override;

    void _accept(Change&) override { msqd += _sqd; }

    void _reject(Change&) override { msqd += 0; }

  protected:
    using Move::spc;
    QuadrantJump(Space& spc, const std::string& name, const std::string& cite);

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

#ifdef ENABLE_MPI

class GibbsEnsembleHelper
{
  public:
    using VectorOfMolIds = std::vector<MoleculeData::index_type>;
    const MPI::Controller& mpi;
    int partner_rank = -1;       
    double total_volume = 0;     
    int total_num_particles = 0; 
    VectorOfMolIds molids;       
    std::pair<int, int>
    currentNumParticles(const Space& spc) const; 
    std::pair<double, double>
    currentVolumes(const Space& spc) const; 
    double exchange(double value) const;    
    GibbsEnsembleHelper(const Space& spc, const MPI::Controller& mpi, const VectorOfMolIds& molids);
};

class GibbsVolumeMove : public VolumeMove
{
  private:
    MPI::Controller& mpi;
    std::unique_ptr<GibbsEnsembleHelper> gibbs;
    bool direct_volume_displacement =
        true; 
    void setNewVolume() override;
    void _from_json(const json& j) override;
    bool volumeTooExtreme() const; 

  protected:
    void _move(Change& change) override;

  public:
    GibbsVolumeMove(Space& spc, MPI::Controller& mpi);
    double bias(Change& change, double old_energy, double new_energy) override;
};

class GibbsMatterMove : public Move
{
  private:
    bool insert;                    
    MoleculeData::index_type molid; 
    MPI::Controller& mpi;
    std::unique_ptr<GibbsEnsembleHelper> gibbs;
    std::unique_ptr<Speciation::GroupDeActivator> molecule_bouncer;
    void _from_json(const json& j) override;
    void _to_json(json& j) const override;

  protected:
    void _move(Change& change) override;

  public:
    GibbsMatterMove(Space& spc, MPI::Controller& mpi);
    double bias(Change& change, double old_energy, double new_energy) override;
};

class ParallelTempering : public Move
{
  private:
    const MPI::Controller& mpi;
    MPI::ExchangeParticles exchange_particles; 
    std::unique_ptr<MPI::Partner> partner;     
    Geometry::VolumeMethod volume_scaling_method =
        Geometry::VolumeMethod::ISOTROPIC;                               
    std::map<MPI::Partner::PartnerPair, Average<double>> acceptance_map; 

    Random slump; // static instance of Random (shared for all in ParallelTempering)
    void _to_json(json& j) const override;
    void _from_json(const json& j) override;
    void _move(Change& change) override;
    void _accept(Change& change) override;
    void _reject(Change& change) override;
    double bias(Change& change, double uold,
                double unew) override;           
    double exchangeEnergy(double energy_change); 
    void exchangeState(Change& change);          
    void exchangeGroupSizes(Space::GroupVector& groups, int partner_rank);

    std::string filename;                 
    std::unique_ptr<std::ostream> stream; 
    std::optional<int> exchange;          
    void writeToFileStream() const;       

  public:
    explicit ParallelTempering(Space& spc, const MPI::Controller& mpi);
};

#endif

[[maybe_unused]] std::unique_ptr<Move> createMove(const std::string& name, const json& properties,
                                                  Space& spc, Energy::Hamiltonian& hamiltonian);

class MoveCollection
{
  private:
    unsigned int number_of_moves_per_sweep; 
    BasePointerVector<Move> moves;          
    std::vector<double> repeats;            
    std::discrete_distribution<unsigned int> distribution; 
    using move_iterator = decltype(moves.vec)::iterator;   
    move_iterator sample(); 

  public:
    MoveCollection(const json& list_of_moves, Space& spc, Energy::Hamiltonian& hamiltonian,
                   Space& old_spc);
    void addMove(std::shared_ptr<Move>&& move); 
    [[maybe_unused]] [[nodiscard]] const BasePointerVector<Move>&
    getMoves() const;                                               
    friend void to_json(json& j, const MoveCollection& propagator); 

    auto repeatedStochasticMoves()
    {
        auto is_valid_and_stochastic = [&](auto move) {
            return move < moves.end() && (*move)->isStochastic();
        };
        using namespace ranges::cpp20;
        return views::iota(0U, number_of_moves_per_sweep) |
               views::transform([&]([[maybe_unused]] auto count) { return sample(); }) |
               views::filter(is_valid_and_stochastic) |
               ranges::views::indirect; // dereference iterator
    }

    auto constantIntervalMoves(const unsigned int sweep_number = 1)
    {
        auto is_static_and_time_to_sample = [&, sweep_number](const auto& move) {
            return (!move->isStochastic()) && (sweep_number % move->sweep_interval == 0);
        };
        return moves | std::views::filter(is_static_and_time_to_sample);
    }
};

void to_json(json& j, const MoveCollection& propagator);

} // namespace move
} // namespace Faunus