geometry.h

#pragma once

#include "units.h"
#include "core.h"
#include "particle.h"
#include "tensor.h"
#include <concepts>
#include <Eigen/Geometry>
#include <spdlog/spdlog.h>
#include <range/v3/view/cartesian_product.hpp>
#include <range/v3/view/zip.hpp>
#include <utility>

namespace Faunus {

class Random;

namespace Geometry {

// Note: BoundaryFunction and DistanceFunction are defined in core.h

enum class Variant
{
    CUBOID = 0,
    SPHERE,
    CYLINDER,
    SLIT,
    HEXAGONAL,
    OCTAHEDRON,
    HYPERSPHERE2D
};

enum class VolumeMethod
{
    ISOTROPIC,
    ISOCHORIC,
    XY,
    Z,
    INVALID
};

NLOHMANN_JSON_SERIALIZE_ENUM(VolumeMethod, {{VolumeMethod::INVALID, nullptr},
                                            {VolumeMethod::ISOTROPIC, "isotropic"},
                                            {VolumeMethod::ISOCHORIC, "isochoric"},
                                            {VolumeMethod::XY, "xy"},
                                            {VolumeMethod::Z, "z"}})

enum class Coordinates
{
    ORTHOGONAL,
    ORTHOHEXAGONAL,
    TRUNC_OCTAHEDRAL,
    NON3D
};
enum class Boundary : int
{
    FIXED = 0,
    PERIODIC = 1
};

class BoundaryCondition
{
  public:
    using BoundaryXYZ = Eigen::Matrix<Boundary, 3, 1>;
    // typedef std::pair<std::string, BoundaryXYZ> BoundaryName;
    // static const std::map<std::string, BoundaryXYZ> names; //!< boundary names

    Coordinates coordinates;
    BoundaryXYZ direction;

    [[nodiscard]] Eigen::Matrix<bool, 3, 1> isPeriodic() const;

    explicit BoundaryCondition(Coordinates coordinates = Coordinates::ORTHOGONAL,
                               BoundaryXYZ boundary = {Boundary::FIXED, Boundary::FIXED,
                                                       Boundary::FIXED})
        : coordinates(coordinates)
        , direction(std::move(boundary)) {};
};

struct GeometryBase
{
    virtual Point setVolume(double, VolumeMethod = VolumeMethod::ISOTROPIC) = 0; 
    [[nodiscard]] virtual double getVolume(int = 3) const = 0;                   
    virtual void boundary(Point&) const = 0;                      
    [[nodiscard]] virtual bool collision(const Point&) const = 0; 
    virtual void randompos(Point&, Random&) const = 0;            
    [[nodiscard]] virtual Point vdist(const Point& a,
                                      const Point& b) const = 0; 
    [[nodiscard]] virtual Point getLength() const = 0;           
    virtual ~GeometryBase();
    virtual void to_json(json& j) const = 0;
    virtual void from_json(const json& j) = 0;

    [[nodiscard]] inline BoundaryFunction getBoundaryFunc() const
    {
        return [this](Point& i) { boundary(i); };
    } 

    [[nodiscard]] inline DistanceFunction getDistanceFunc() const
    {
        return [this](const Point& i, const Point& j) { return vdist(i, j); };
    } 

  protected:
    template <typename T = double> [[nodiscard]] inline int anint(T x) const
    {
        return int(x > 0.0 ? x + 0.5 : x - 0.5);
    } 

}; 

class GeometryImplementation : public GeometryBase
{
  public:
    BoundaryCondition boundary_conditions;

    ~GeometryImplementation() override;

    [[nodiscard]] virtual std::unique_ptr<GeometryImplementation> clone() const = 0;
};

class Cuboid : public GeometryImplementation
{
  protected:
    Point box, box_half, box_inv;

  public:
    Point getLength() const override;
    double getVolume(int dim = 3) const final; // finalized to help the compiler with inlining
    void setLength(const Point& len);          // todo shall be protected
    Point setVolume(double volume, VolumeMethod method = VolumeMethod::ISOTROPIC) override;
    Point vdist(const Point& a, const Point& b) const override; 
    void boundary(Point& a) const override;
    bool collision(const Point& a) const override;
    void randompos(Point& m, Random& rand) const override;
    void from_json(const json& j) override;
    void to_json(json& j) const override;
    explicit Cuboid(const Point& side_length);
    Cuboid();

    [[nodiscard]] std::unique_ptr<GeometryImplementation>
    clone() const override; 
};

class Slit : public Cuboid
{
    using Tbase = Cuboid;

  public:
    explicit Slit(const Point& p);
    Slit(double x, double y, double z);
    explicit Slit(double x = 0.0);

    [[nodiscard]] std::unique_ptr<GeometryImplementation>
    clone() const override; 
};

class Sphere : public GeometryImplementation
{
  protected:
    double radius;

  public:
    Point getLength() const override;
    double getVolume(int dim = 3) const override;
    Point setVolume(double volume, VolumeMethod method = VolumeMethod::ISOTROPIC) override;
    Point vdist(const Point& a, const Point& b) const override; 

    inline static double sqdist(const Point& a, const Point& b) { return (a - b).squaredNorm(); };

    void boundary(Point& a) const override;
    bool collision(const Point& point) const override;
    void randompos(Point& m, Random& rand) const override;
    void from_json(const json& j) override;
    void to_json(json& j) const override;
    explicit Sphere(double radius = 0.0);
    double getRadius() const;

    std::unique_ptr<GeometryImplementation>
    clone() const override; 
};

class Hypersphere2d : public Sphere
{
  public:
    Point vdist(const Point& a, const Point& b) const override;
    bool collision(const Point& a) const override;
    void randompos(Point& m, Random& rand) const override;
    explicit Hypersphere2d(double radius = 0.0);

    std::unique_ptr<GeometryImplementation>
    clone() const override; 
};

class Cylinder : public GeometryImplementation
{
  protected:
    double radius, height;

  public:
    Point getLength() const override;
    double getVolume(int dim = 3) const override;
    Point setVolume(double volume, VolumeMethod method = VolumeMethod::ISOTROPIC) override;
    Point vdist(const Point& a, const Point& b) const override;
    void boundary(Point& a) const override;
    bool collision(const Point& a) const override;
    void randompos(Point& m, Random& rand) const override;
    void from_json(const json& j) override;
    void to_json(json& j) const override;
    explicit Cylinder(double radius = 0.0, double height = 0.0);

    [[nodiscard]] std::unique_ptr<GeometryImplementation>
    clone() const override; 
};

class HexagonalPrism : public GeometryImplementation
{
    static const Eigen::Matrix3d rhombic2cartesian;
    static const Eigen::Matrix3d cartesian2rhombic;

    Point box; 
    void set_box(double side, double height);

  public:
    Point getLength() const override;
    double getVolume(int dim = 3) const override;
    Point setVolume(double volume, VolumeMethod method = VolumeMethod::ISOTROPIC) override;
    Point vdist(const Point& a, const Point& b) const override;
    void boundary(Point& a) const override;
    bool collision(const Point& a) const override;
    void randompos(Point& m, Random& rand) const override;
    void from_json(const json& j) override;
    void to_json(json& j) const override;
    explicit HexagonalPrism(double side = 0.0, double height = 0.0);

    [[nodiscard]] std::unique_ptr<GeometryImplementation>
    clone() const override; 

    [[nodiscard]] double innerRadius() const; 
    [[nodiscard]] double outerRadius() const; 
    [[nodiscard]] double height() const;      
};

class TruncatedOctahedron : public GeometryImplementation
{
    double side;

  public:
    Point getLength() const override;
    double getVolume(int dim = 3) const override;
    Point setVolume(double volume, VolumeMethod method = VolumeMethod::ISOTROPIC) override;
    Point vdist(const Point& a, const Point& b) const override;
    void boundary(Point& a) const override;
    bool collision(const Point& a) const override;
    void randompos(Point& pos, Random& rand) const override;
    void from_json(const json& j) override;
    void to_json(json& j) const override;
    explicit TruncatedOctahedron(double side = 0.0);

    [[nodiscard]] std::unique_ptr<GeometryImplementation>
    clone() const override; 
};

class Chameleon : public GeometryBase
{
  private:
    Point len_or_zero = {0, 0, 0}; 
    Point len, len_half,
        len_inv; 
    std::unique_ptr<GeometryImplementation> geometry =
        nullptr;       
    Variant _type;     
    std::string _name; 
    void
    makeGeometry(const Variant type =
                     Variant::CUBOID); 
    void _setLength(const Point& l);

  public:
    const Variant& type = _type;     
    const std::string& name = _name; 
    double getVolume(int dim = 3) const override;
    Point setVolume(double, VolumeMethod = VolumeMethod::ISOTROPIC) override;
    Point getLength() const override; 
    // setLength() needed only for Move::ReplayMove (stems from IO::XTCReader).
    void setLength(const Point&);                           
    void boundary(Point&) const override;                   
    Point vdist(const Point&, const Point&) const override; 
    double sqdist(const Point&,
                  const Point&) const; 
    void randompos(Point&, Random&) const override;
    bool collision(const Point&) const override;
    void from_json(const json&) override;
    void to_json(json& j) const override;

    const BoundaryCondition& boundaryConditions() const; 

    static const std::map<std::string, Variant> names; 
    typedef std::pair<std::string, Variant> VariantName;

    static VariantName variantName(const std::string& name);

    static VariantName variantName(const json& j);

    explicit Chameleon(const Variant type = Variant::CUBOID);
    Chameleon(const GeometryImplementation& geo, Variant type);

    Chameleon(const Chameleon& geo);

    Chameleon& operator=(const Chameleon& geo);

    [[nodiscard]] std::shared_ptr<GeometryImplementation> asSimpleGeometry() const;
};

inline void Chameleon::randompos(Point& m, Random& rand) const
{
    assert(geometry);
    geometry->randompos(m, rand);
}

inline bool Chameleon::collision(const Point& a) const
{
    assert(geometry);
    return geometry->collision(a);
}

inline void Chameleon::boundary(Point& a) const
{
    const auto& boundary_conditions = geometry->boundary_conditions;
    if (boundary_conditions.coordinates == Coordinates::ORTHOGONAL) {
        if (boundary_conditions.direction.x() == Boundary::PERIODIC) {
            if (std::fabs(a.x()) > len_half.x())
                a.x() -= len.x() * anint(a.x() * len_inv.x());
        }
        if (boundary_conditions.direction.y() == Boundary::PERIODIC) {
            if (std::fabs(a.y()) > len_half.y())
                a.y() -= len.y() * anint(a.y() * len_inv.y());
        }
        if (boundary_conditions.direction.z() == Boundary::PERIODIC) {
            if (std::fabs(a.z()) > len_half.z())
                a.z() -= len.z() * anint(a.z() * len_inv.z());
        }
    }
    else {
        geometry->boundary(a);
    }
}

inline Point Chameleon::vdist(const Point& a, const Point& b) const
{
    Point distance;
    const auto& boundary_conditions = geometry->boundary_conditions;
    if (boundary_conditions.coordinates == Coordinates::ORTHOGONAL) {
        distance = a - b;
        if (boundary_conditions.direction.x() == Boundary::PERIODIC) {
            if (distance.x() > len_half.x())
                distance.x() -= len.x();
            else if (distance.x() < -len_half.x())
                distance.x() += len.x();
        }
        if (boundary_conditions.direction.y() == Boundary::PERIODIC) {
            if (distance.y() > len_half.y())
                distance.y() -= len.y();
            else if (distance.y() < -len_half.y())
                distance.y() += len.y();
        }
        if (boundary_conditions.direction.z() == Boundary::PERIODIC) {
            if (distance.z() > len_half.z())
                distance.z() -= len.z();
            else if (distance.z() < -len_half.z())
                distance.z() += len.z();
        }
    }
    else {
        distance = geometry->vdist(a, b);
    }
    return distance;
}

inline double Chameleon::sqdist(const Point& a, const Point& b) const
{
    if (geometry->boundary_conditions.coordinates == Coordinates::ORTHOGONAL) {
        if constexpr (true) {
            Point d((a - b).cwiseAbs());
            return (d - (d.array() > len_half.array())
                            .cast<double>()
                            .matrix()
                            .cwiseProduct(len_or_zero))
                .squaredNorm();
        }
        else { // more readable alternative(?), nearly same speed
            Point d(a - b);
            for (int i = 0; i < 3; ++i) {
                d[i] = std::fabs(d[i]);
                d[i] = d[i] -
                       len_or_zero[i] *
                           static_cast<double>(d[i] > len_half[i]); // casting faster than branching
            }
            return d[0] * d[0] + d[1] * d[1] + d[2] * d[2];
        }
    }
    else {
        return geometry->vdist(a, b).squaredNorm();
    }
}

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

enum class weight
{
    MASS,
    CHARGE,
    GEOMETRIC
};

template <std::ranges::range Positions, std::ranges::range Weights>
Point weightedCenter(
    const Positions& positions, const Weights& weights,
    Geometry::BoundaryFunction boundary = [](auto&) {}, const Point& shift = Point::Zero())
{
    double weight_sum = 0.0;
    Point center(0.0, 0.0, 0.0);
    for (const auto& [position, weight] : ranges::views::zip(positions, weights)) {
        Point shifted_position = position + shift;
        boundary(shifted_position);
        center += weight * shifted_position;
        weight_sum += weight;
    }
    if (std::fabs(weight_sum) > pc::epsilon_dbl) {
        center = center / weight_sum - shift; // translate back
        boundary(center);
        return center;
    }
    else {
        faunus_logger->warn("warning: sum of weights is zero! setting center to (0,0,0)");
        return Point::Zero();
    }
}

template <RequireParticleIterator iterator> //, typename weightFunc>
Point weightedCenter(iterator begin, iterator end, BoundaryFunction boundary,
                     std::function<double(const Particle&)> weight_function,
                     const Point& shift = Point::Zero())
{
    namespace rv = std::views;
    auto particles = std::ranges::subrange(begin, end);
    auto positions = particles | rv::transform(&Particle::pos);
    auto weights = particles | rv::transform(weight_function);
    return weightedCenter(positions, weights, boundary, shift);
} 

template <RequireParticleIterator iterator>
Point massCenter(
    iterator begin, iterator end, BoundaryFunction apply_boundary = [](Point&) {},
    const Point& shift = {0.0, 0.0, 0.0})
{
    auto particle_mass = [](const auto& particle) -> double { return particle.traits().mw; };
    return weightedCenter(begin, end, apply_boundary, particle_mass, shift);
}

template <RequireParticleIterator iterator>
void translate(
    iterator begin, iterator end, const Point& displacement,
    BoundaryFunction apply_boundary = [](auto&) {})
{
    std::for_each(begin, end, [&](auto& particle) {
        particle.pos += displacement;
        apply_boundary(particle.pos);
    });
}

template <RequireParticleIterator iterator>
void translateToOrigin(iterator begin, iterator end, BoundaryFunction apply_boundary = [](auto&) {})
{
    Point cm = massCenter(begin, end, apply_boundary);
    translate(begin, end, -cm, apply_boundary);
}

template <RequireParticleIterator iterator>
void rotate(
    iterator begin, iterator end, const Eigen::Quaterniond& quaternion,
    BoundaryFunction apply_boundary = [](auto&) {}, const Point& shift = Point::Zero())
{
    const auto rotation_matrix = quaternion.toRotationMatrix(); // rotation matrix
    std::for_each(begin, end, [&](auto& particle) {
        particle.rotate(quaternion, rotation_matrix); // rotate internal coordinates
        particle.pos += shift;
        apply_boundary(particle.pos);
        particle.pos = (quaternion * particle.pos) - shift; // rotate positions
        apply_boundary(particle.pos);
    });
} 

template <class Tspace, class GroupIndex>
Point trigoCom(const Tspace& spc, const GroupIndex& indices,
               const std::vector<int>& dir = {0, 1, 2})
{
    if (dir.empty() || dir.size() > 3) {
        throw std::out_of_range("invalid directions");
    }
    namespace rv = std::views;
    using std::views::join;
    auto positions = indices | rv::transform([&](auto i) { return spc.groups.at(i); }) | join |
                     rv::transform(&Particle::pos);
    Point xhi(0, 0, 0);
    Point zeta(0, 0, 0);
    Point theta(0, 0, 0);
    Point com(0, 0, 0);
    for (auto k : dir) {
        const auto q = 2.0 * pc::pi / spc.geometry.getLength()[k];
        int cnt = 0;
        std::for_each(positions.begin(), positions.end(), [&](auto& position) {
            theta[k] = position[k] * q;
            zeta[k] += std::sin(theta[k]);
            xhi[k] += std::cos(theta[k]);
            cnt++;
        });
        theta[k] =
            std::atan2(-zeta[k] / static_cast<double>(cnt), -xhi[k] / static_cast<double>(cnt)) +
            pc::pi;
        com[k] = spc.geometry.getLength()[k] * theta[k] / (2.0 * pc::pi);
    }
    spc.geometry.boundary(com); // is this really needed?
    return com;
}

template <RequirePointIterator position_iterator, std::forward_iterator mass_iterator>
Tensor gyration(
    position_iterator begin, position_iterator end, mass_iterator mass, const Point& mass_center,
    const BoundaryFunction boundary = [](auto&) {})
{
    Tensor S = Tensor::Zero();
    double total_mass = 0.0;
    std::for_each(begin, end, [&](const Point& position) {
        Point r = position - mass_center; // get rid...
        boundary(r);                      // ...of PBC (if any)
        S += (*mass) * r * r.transpose();
        total_mass += (*mass);
        std::advance(mass, 1);
    });
    if (total_mass > pc::epsilon_dbl) {
        return S / total_mass;
    }
    throw std::runtime_error("gyration tensor: total mass must be positive");
}

template <RequireParticleIterator iterator>
Tensor gyration(
    iterator begin, iterator end, const Point& mass_center,
    const BoundaryFunction boundary = [](auto&) {})
{
    namespace rv = std::views;
    auto particles = std::ranges::subrange(begin, end);
    auto positions = particles | rv::transform(&Particle::pos);
    auto masses = particles | rv::transform(&Particle::traits) | rv::transform(&AtomData::mw);
    return gyration(positions.begin(), positions.end(), masses.begin(), mass_center, boundary);
}

struct ShapeDescriptors
{
    double gyration_radius_squared = 0.0;
    double asphericity = 0.0;
    double acylindricity = 0.0;
    double relative_shape_anisotropy =
        0.0; 
    ShapeDescriptors() = default;
    ShapeDescriptors(const Tensor& gyration_tensor); 
    ShapeDescriptors&
    operator+=(const ShapeDescriptors& other); 
    ShapeDescriptors operator*(const double scale) const; 
};

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

template <RequireParticleIterator iterator>
Tensor inertia(
    iterator begin, iterator end, const Point origin = Point::Zero(),
    const BoundaryFunction boundary = [](auto&) {})
{
    Tensor I = Tensor::Zero();
    std::for_each(begin, end, [&](const Particle& particle) {
        Point t = particle.pos - origin;
        boundary(t);
        I += particle.traits().mw * (t.squaredNorm() * Tensor::Identity() - t * t.transpose());
    });
    return I;
}

template <std::forward_iterator InputIt1, std::forward_iterator InputIt2, typename BinaryOperation>
double rootMeanSquareDeviation(InputIt1 begin, InputIt1 end, InputIt2 d_begin,
                               BinaryOperation diff_squared_func)
{
    assert(std::distance(begin, end) > 0);
    double sq_sum = 0;
    for (InputIt1 i = begin; i != end; ++i) {
        sq_sum += diff_squared_func(*i, *d_begin);
        ++d_begin;
    }
    return std::sqrt(sq_sum / std::distance(begin, end));
}

ParticleVector mapParticlesOnSphere(const ParticleVector&);

std::pair<Cuboid, ParticleVector> hexagonalPrismToCuboid(const HexagonalPrism& hexagon,
                                                         const RequireParticles auto& particles)
{
    Cuboid cuboid({2.0 * hexagon.innerRadius(), 3.0 * hexagon.outerRadius(), hexagon.height()});
    ParticleVector cuboid_particles;
    cuboid_particles.reserve(2 * std::distance(particles.begin(), particles.end()));
    std::copy(particles.begin(), particles.end(),
              std::back_inserter(cuboid_particles)); // add central hexagon

    std::transform(particles.begin(), particles.end(), std::back_inserter(cuboid_particles),
                   [&](auto particle) {
                       particle.pos.x() +=
                           hexagon.innerRadius() * (particle.pos.x() > 0.0 ? -1.0 : 1.0);
                       particle.pos.y() +=
                           hexagon.outerRadius() * (particle.pos.y() > 0.0 ? -1.5 : 1.5);
                       assert(cuboid.collision(particle.pos) == false);
                       return particle;
                   }); // add the four corners; i.e. one extra, split hexagon
    assert(std::fabs(cuboid.getVolume() - 2.0 * hexagon.getVolume()) <= pc::epsilon_dbl);
    return {cuboid, cuboid_particles};
}

class TwobodyAngles
{
    static std::vector<Point> fibonacciSphere(size_t);

  public:
    std::vector<Eigen::Quaterniond> quaternions_1; 
    std::vector<Eigen::Quaterniond> quaternions_2; 
    std::vector<Eigen::Quaterniond> dihedrals;     

    TwobodyAngles() = default;
    TwobodyAngles(double angle_resolution);

    auto quaternionPairs() const
    {
        namespace rv = ranges::views;
        auto product = [](const auto& pair) -> Eigen::Quaterniond {
            const auto& [q2, q_dihedral] = pair;
            return q_dihedral * q2; // noncommutative
        };
        auto second_body_quaternions =
            rv::cartesian_product(quaternions_2, dihedrals) | std::views::transform(product);
        return rv::cartesian_product(quaternions_1, second_body_quaternions);
    }

    size_t
    size() const; 
};

class TwobodyAnglesState : public TwobodyAngles
{
  private:
    using QuaternionIter = std::vector<Eigen::Quaterniond>::const_iterator;
    QuaternionIter q_euler1;
    QuaternionIter q_euler2;
    QuaternionIter q_dihedral;

  public:
    TwobodyAnglesState() = default;
    TwobodyAnglesState(double angle_resolution);
    TwobodyAnglesState& advance(); 
    std::optional<std::pair<Eigen::Quaterniond, Eigen::Quaterniond>>
    get(); 
};

} // namespace Geometry
} // namespace Faunus