/home/js850/projects/pele/source/pele/aatopology.h

#ifndef _PELE_AATOPOLOGY_H_
#define _PELE_AATOPOLOGY_H_

#include <string>
#include <list>
#include <cmath>
#include <stdexcept>

#include "pele/array.h"
#include "pele/rotations.h"
#include "pele/base_potential.h"
#include "pele/vecn.h"
#include "pele/lowest_eig_potential.h"
#include "pele/matrix.h"
#include "pele/distance.h"

namespace pele{




class CoordsAdaptor {
    size_t _nrigid;  
    size_t _natoms;  
    pele::Array<double> _coords;
    static const size_t _ndim = 3;

public:
    CoordsAdaptor(size_t nrigid, size_t natoms, pele::Array<double> coords)
        : _nrigid(nrigid),
          _natoms(natoms),
          _coords(coords)
    {
        if (coords.size() != (6*_nrigid + 3*_natoms)) {
            throw std::invalid_argument(std::string("coords has the wrong size ") + std::to_string(coords.size()));
        }
        if (natoms != 0)
            throw std::runtime_error("coords adaptor doesn't support free atoms yet");
    }

    pele::Array<double> get_coords() { return _coords; }

    pele::Array<double> get_rb_positions()
    {
        if (_nrigid == 0) {
            // return empty array
            return pele::Array<double>();
        }
        return _coords.view(0, 3*_nrigid);
    }

    pele::Array<double> get_rb_position(size_t isite)
    {
        if (_nrigid == 0) {
            // return empty array
            return pele::Array<double>();
        }
        if (isite > _nrigid) {
            throw std::invalid_argument("isite must be less than nrigid");
        }
        size_t const istart = 3*isite;
        return _coords.view(istart, istart+3);
    }

    pele::Array<double> get_rb_rotations()
    {
        if (_nrigid == 0) {
            // return empty array
            return pele::Array<double>();
        }
        return _coords.view(3*_nrigid, 6*_nrigid);
    }

   pele::Array<double> get_rb_rotation(size_t isite)
    {
        if (_nrigid == 0) {
            // return empty array
            return pele::Array<double>();
        }
        if (isite > _nrigid) {
            throw std::invalid_argument("isite must be less than nrigid");
        }
        size_t const istart = 3*_nrigid + 3*isite;
        return _coords.view(istart, istart+3);
    }

    pele::Array<double> get_atom_positions()
    {
        if (_natoms == 0) {
            // return empty array
            return pele::Array<double>();
        }

        return _coords.view(6*_nrigid, 6*_nrigid + 3*_natoms);
    }

};

// forward definition of RBTopology needed for TrasnformAACluster
class RBTopology;

class TransformPolicy {
public:
//    void translate(self, X, d) {
//        ''' translate the coordinates '''
//    }
    virtual ~TransformPolicy() {}

    virtual void rotate(pele::Array<double> x, pele::MatrixNM<3,3> const & mx) = 0;

//    def can_invert(self):
//        ''' returns True or False if an inversion can be performed'''
//
//    def invert(self, X):
//        ''' perform an inversion at the origin '''
//
//    def permute(self, X, perm):
//        ''' returns the permuted coordinates '''

};

class TransformAACluster : public TransformPolicy {
public:
    pele::RBTopology * m_topology;
    TransformAACluster(pele::RBTopology * topology)
        : m_topology(topology)
    { }
    virtual ~TransformAACluster() {}

    void rotate(pele::Array<double> x, pele::MatrixNM<3,3> const & mx);
//    inline void rotate(pele::Array<double> x, pele::Array<double> mx)
//    {
//        return rotate(x, pele::MatrixNM<3,3>(mx));
//    }
};

class MeasureAngleAxisCluster {
public:
    pele::RBTopology * m_topology;
    MeasureAngleAxisCluster(pele::RBTopology * topology)
        : m_topology(topology)
    { }

    void align(pele::Array<double> const x1, pele::Array<double> x2);


};


class RigidFragment {
    static const size_t _ndim = 3;
    pele::Array<double> _atom_positions;
    pele::MatrixAdapter<double> _atom_positions_matrix;
    size_t _natoms;

    std::shared_ptr<DistanceInterface> m_distance_function;

    double m_M; // total mass of the angle axis site
    double m_W; // sum of all weights
    pele::VecN<3> m_cog; // center of gravity
    pele::MatrixNM<3,3> m_S; // weighted tensor of gyration S_ij = \sum m_i x_i x_j
    pele::MatrixNM<3,3> m_inversion; // matrix that applies the appropriate inversion
    bool m_can_invert;

    // a list of rotations that leave the rigid body unchanged.
    std::vector<pele::MatrixNM<3,3> > m_symmetry_rotations;


public:
    RigidFragment(pele::Array<double> atom_positions,
            Array<double> cog,
            double M,
            double W,
            Array<double> S,
            Array<double> inversion, bool can_invert,
            std::shared_ptr<DistanceInterface> distance_function
            )
    : _atom_positions(atom_positions.copy()),
      _atom_positions_matrix(_atom_positions, _ndim),
      _natoms(_atom_positions.size() / _ndim),
      m_distance_function(distance_function),
      m_M(M),
      m_W(W),
      m_cog(cog),
      m_S(S),
      m_inversion(inversion),
      m_can_invert(can_invert)
    {
        if (_atom_positions.size() == 0 ) {
            throw std::invalid_argument("the atom positions must not have zero size");
        }
        if (_atom_positions.size() != _natoms * _ndim ) {
            throw std::invalid_argument("the length of atom_positions must be divisible by 3");
        }
    }

    inline size_t natoms() const { return _natoms; }

    inline void add_symmetry_rotation(pele::Array<double> R)
    {
        m_symmetry_rotations.push_back(R);
    }

    inline std::vector<pele::MatrixNM<3,3> > const & get_symmetry_rotations() const
    {
        return m_symmetry_rotations;
    }

    pele::Array<double> to_atomistic(pele::Array<double> const com,
            pele::VecN<3> const & p);

    void transform_grad(
            pele::VecN<3> const & p,
            pele::Array<double> const g,
            pele::VecN<3> & g_com,
            pele::VecN<3> & g_rot
            );

    void transform_grad(
            pele::Array<double> const p,
            pele::Array<double> const g,
            pele::Array<double> g_com,
            pele::Array<double> g_rot
            )
    {
        pele::VecN<3> p_vec = p;
        pele::VecN<3> g_com_vec = g_com;
        pele::VecN<3> g_rot_vec = g_rot;
        transform_grad(p_vec, g, g_com_vec, g_rot_vec);
        for (size_t i = 0; i<3; ++i) {
            g_com[i] = g_com_vec[i];
            g_rot[i] = g_rot_vec[i];
        }

    }

    inline pele::VecN<3> get_smallest_rij(pele::VecN<3> const & com1, pele::VecN<3> const & com2) const
    {
        pele::VecN<3> distance;
        m_distance_function->get_rij(distance.data(), com2.data(), com1.data());
        return distance;
    }

    double distance_squared(pele::VecN<3> const & com1, pele::VecN<3> const & p1,
            pele::VecN<3> const & com2, pele::VecN<3> const & p2) const;

    void distance_squared_grad(pele::VecN<3> const & com1, pele::VecN<3> const & p1,
            pele::VecN<3> const & com2, pele::VecN<3> const & p2,
            VecN<3> & g_M, VecN<3> & g_P
            ) const;
};

// * Angle axis topology
// *
// * An angle axis topology stores all topology information for an angle axis
// * system. The AATopology is composed of several angle axis sites,
// * which describe the shape of the angle axis site and each site carries a
// * position and orientation. Therefore, the length of the coordinate array
// * must be 6*number_of_sites.
// */
//class AATopology {
//};

class RBTopology {
    std::vector<RigidFragment> _sites;
    size_t _natoms_total;

public:
    RBTopology()
        : _natoms_total(0)
    {}

    void add_site(RigidFragment const & site)
    {
        _sites.push_back(site);
        _natoms_total += site.natoms();
    }

    std::vector<RigidFragment> const & get_sites() const { return _sites; };

    size_t nrigid() const { return _sites.size(); }

    size_t natoms_total() const { return _natoms_total; }

    size_t number_of_non_rigid_atoms() const { return 0; }

    CoordsAdaptor get_coords_adaptor(pele::Array<double> x) const
    {
        return CoordsAdaptor(nrigid(), number_of_non_rigid_atoms(), x);
    }

    Array<double> to_atomistic(Array<double> rbcoords);

    void transform_gradient(pele::Array<double> rbcoords,
            pele::Array<double> grad, pele::Array<double> rbgrad);

    pele::VecN<3> align_angle_axis_vectors(pele::VecN<3> const & p1,
            pele::VecN<3> const & p2in);

    void align_all_angle_axis_vectors(pele::Array<double> x1,
            pele::Array<double> x2);

    void align_path(std::list<pele::Array<double> > path);

    void get_zero_modes(pele::Array<double> const x,
            std::vector<pele::Array<double> > & zev)
    {
        auto ca = get_coords_adaptor(x);
        pele::Array<double> v(x.size(), 0);
        auto cv = get_coords_adaptor(v);

        // get the zero eigenvectors corresponding to translation
        std::vector<pele::Array<double> > zev_t;
        pele::zero_modes_translational(zev_t, nrigid(), 3);

        for (auto const & v : zev_t) {
            cv.get_rb_positions().assign(v);
            zev.push_back(cv.get_coords().copy());
        }

        // get the zero eigenvectors corresponding to rotation
        TransformAACluster transform(this);
        double d = 1e-5;
        pele::VecN<3> v3;
        pele::Array<double> delta(x.size());

        // do rotations around the x y and z axes
        for (size_t i = 0; i < 3; ++i) {
            delta.assign(x);
            v3.assign(0);
            v3[i] = d;
            transform.rotate(delta, pele::aa_to_rot_mat(v3));
            align_all_angle_axis_vectors(x, delta);
            delta -= x;
            delta /= norm(delta);
            zev.push_back(delta.copy());
        }
    }

    double distance_squared(pele::Array<double> const x1, pele::Array<double> const x2) const
    {
        double d_sq = 0;
        auto ca1 = get_coords_adaptor(x1);
        auto ca2 = get_coords_adaptor(x2);
        for (size_t isite = 0; isite < nrigid(); ++isite) {
            d_sq += _sites[isite].distance_squared(
                    ca1.get_rb_position(isite),
                    ca1.get_rb_rotation(isite),
                    ca2.get_rb_position(isite),
                    ca2.get_rb_rotation(isite)
                );
        }
        return d_sq;
    }

    void distance_squared_grad(pele::Array<double> const x1, pele::Array<double> const x2,
            pele::Array<double> grad
            ) const
    {
        if (grad.size() != x1.size()) {
            throw std::runtime_error("grad has the wrong size");
        }
        grad.assign(0);
        auto ca1 = get_coords_adaptor(x1);
        auto ca2 = get_coords_adaptor(x2);
        auto ca_spring = get_coords_adaptor(grad);

        // first distance for sites only
        for (size_t isite=0; isite<nrigid(); ++isite) {
            pele::VecN<3> g_M, g_P;
            _sites[isite].distance_squared_grad(
                    ca1.get_rb_position(isite),
                    ca1.get_rb_rotation(isite),
                    ca2.get_rb_position(isite),
                    ca2.get_rb_rotation(isite),
                    g_M, g_P);
            auto spring_com = ca_spring.get_rb_position(isite);
            std::copy(g_M.begin(), g_M.end(), spring_com.begin());
            auto spring_rot = ca_spring.get_rb_rotation(isite);
            std::copy(g_P.begin(), g_P.end(), spring_rot.begin());
        }
    }
};

class RBPotentialWrapper : public BasePotential {
    std::shared_ptr<BasePotential> potential_;
    std::shared_ptr<RBTopology> topology_;
public:

    RBPotentialWrapper(std::shared_ptr<BasePotential> potential,
            std::shared_ptr<RBTopology> top)
        : potential_(potential),
          topology_(top)

    {}

    inline double get_energy(pele::Array<double> rbcoords)
    {
        auto x = topology_->to_atomistic(rbcoords);
        return potential_->get_energy(x);
    }

    inline double get_energy_gradient(pele::Array<double> rbcoords,
            pele::Array<double> rbgrad)
    {
        auto x = topology_->to_atomistic(rbcoords);
        pele::Array<double> grad_atomistic(topology_->natoms_total() * 3);
        double e = potential_->get_energy_gradient(x, grad_atomistic);
        topology_->transform_gradient(rbcoords, grad_atomistic, rbgrad);
        return e;
    }
};

}

#endif