/**
 * @brief Handcrafted flat 3Dtree.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 *
 * @note This module is for reference purpose and much slower than nanoflann library.
 */
#ifndef TBPLAS_BUILDER_MESHGRID_H

#include <list>

#include <eigen3/Eigen/Dense>

#include "base.h"

namespace tbplas::builder {
/**
 * @brief Cubic volume in 3D space for holding points.
 *
 * @note Cube origin, size and unit vectors are handled by MeshGrid class.
 */
class Cube {
public:
    /**
     * @brief Add a new point to the cube.
     *
     * @param[in] idx index of the point
     */
    void add_point(const size_t& idx);

    /**
     * @brief Get the number of points within this cube.
     *
     * @return number of points
     */
    [[nodiscard]] size_t get_num_points() const;

    /**
     * @brief Get the indices of points
     *
     * @return indices of points
     */
    [[nodiscard]] const std::vector<size_t>& get_point_indices() const;

    /**
     * @brief Set the cube index
     * 
     * @param[in] idx cube index in sc representation
     */
    void set_cube_index(const int64_t& idx);

    /**
     * @brief Get the cube index
     * 
     * @return cube index in sc representation
     */
    [[nodiscard]] const int64_t& get_cube_index() const;

private:
    /// @brief Indices of the points
    std::vector<size_t> point_indices_;

    /// @brief Index of the cube in sc representation
    int64_t cube_index_;
};

/**
 * @brief Similar to Term in kdtree.h, but without rn.
 *
 */
class ReducedTerm {
public:
    /// @brief orbital indices
    pair_t pair;

    /// @brief displacement vector in nm
    Eigen::Vector3d rij;

    /// @brief norm of rij in nm
    double distance;

    // Constructor
    explicit ReducedTerm(
        const pair_t& pair,
        const Eigen::Vector3d& rij,
        const double& distance);

    // Operators for sorting
    bool operator<(const ReducedTerm& other) const;
    bool operator>(const ReducedTerm& other) const;
};

/**
 * @brief Partition of 3D space, for boosting KDTree.
 *
 */
class MeshGrid {
public:
    /**
     * @brief Construct a new MeshGrid object
     *
     * @param[in] pos_bra Cartesian coordinates of pos_bra, dimension (3, num_points_bra)
     * @param[in] pos_ket Cartesian coordinates of pos_ket, dimension (3, num_points_ket)
     * @param[in] cube_size cube size in same unit as pos_bra and pos_ket
     *
     * @note Cube size is typically 1-4 times of max_distance.
     */
    explicit MeshGrid(
        const Eigen::Matrix3Xd& pos_bra,
        const Eigen::Matrix3Xd& pos_ket,
        const double& cube_size);

    /**
     * @brief Get the neighbors up to given cutoff distance.
     *
     * @param[in] max_distance cutoff distance
     * @return list of neighbors
     */
    [[nodiscard]] std::list<ReducedTerm>
    get_neighbors(const double& max_distance = 1.0);

private:
    /// @brief Pointer to Cartesian coordinates of pos_bra
    const Eigen::Matrix3Xd* pos_bra_ = nullptr;

    /// @brief Pointer to Cartesian coordinates of pos_ket
    const Eigen::Matrix3Xd* pos_ket_ = nullptr;

    /// @brief Cartesian coordinate of origin
    Eigen::Vector3d origin_ = Eigen::Vector3d::Zero();

    /// @brief Lattice vector of unit volume in ROW-WISE order
    Eigen::Matrix3d unit_vec_ = Eigen::Matrix3d::Identity();

    /// @brief Dimension of mesh grid
    std::vector<int64_t> dim_ = { 1, 1, 1 };

    /// @brief Cubes within this mesh grid
    std::vector<Cube> cubes_bra_;
    std::vector<Cube> cubes_ket_;

    /**
     * @brief Participate 3D space.
     *
     * @param[in] cube_size cube size
     */
    void set_partition(const double& cube_size);

    /**
     * @brief Convert cube index from pc to sc representation.
     *
     * @param[in] ra index along a direction
     * @param[in] rb index along b direction
     * @param[in] rc index along c direction
     * @return cube index in sc representation
     */
    [[nodiscard]] int64_t id_pc2sc(
        const int64_t& ra,
        const int64_t& rb,
        const int64_t& rc) const;

    /**
     * @brief Convert cube index from sc to pc representation.
     *
     * @param[in] id_sc cube index in sc representation
     * @param[out] ra index along a direction
     * @param[out] rb index along b direction
     * @param[out] rc index along c direction
     */
    void id_sc2pc(
        const int64_t& id_sc,
        int64_t& ra,
        int64_t& rb,
        int64_t& rc) const;

    /**
     * @brief Python-compatible version of division j0 / n0
     *
     * @param[in] j0 dividend
     * @param[in] n0 divisor
     * @param[out] j1 = j0 % n0
     * @param[out] n1 = j0 // n0
     */
    static void py_div(
        const int64_t& j0,
        const int64_t& n0,
        int64_t& j1,
        int64_t& n1)
    {
        int64_t j1_new = j0;
        while (j1_new < 0) {
            j1_new += n0;
        }
        while (j1_new >= n0) {
            j1_new -= n0;
        }
        j1 = j1_new;
        n1 = (j0 - j1_new) / n0;
    }

    /**
     * @brief Get the floor of a vector
     *
     * @param[in] vec vector to split
     * @param[out] ra floor on a direction
     * @param[out] rb floor on b direction
     * @param[out] rc floor on c direction
     * @param[in] delta small displacement added to vector to avoid clipping
     */
    static void floor_vector3(
        const Eigen::Vector3d& vec,
        int64_t& ra,
        int64_t& rb,
        int64_t& rc,
        const double& delta = 0.02)
    {
        ra = std::floor(vec(0) + delta);
        rb = std::floor(vec(1) + delta);
        rc = std::floor(vec(2) + delta);
    }

    /**
     * @brief Divide points into cubes.
     *
     * @param[in] pos Cartesian coordinates of points, should be either pos_bra_ or pos_ket_
     * @param[inout] cubes container of cubes, should be either cubes_bra_ or cubes_ket_
     */
    void divide_points(const Eigen::Matrix3Xd* pos, std::vector<Cube>& cubes);

    /**
     * @brief Get neighboring cubes of cube_bra in cubes_ket.
     *
     * @param[in] cube_bra pointer to cube_bra
     * @return pointers to neighboring cubes
     *
     * @note Empty cubes are excluded from the result.
     */
    [[nodiscard]] std::vector<const Cube*>
    get_neighbor_cubes_ket(const Cube* cube_bra) const;

    /**
     * @brief Add neighbors from cube_bra to cube_ket.
     * 
     * @param[in] cube_bra pointer to cube_bra
     * @param[in] cube_ket pointer to cube_ket
     * @param[in] max_distance cutoff distance
     * @param[inout] neighbors container of neighbors
     */
    void add_neighbors(
        const Cube* cube_bra,
        const Cube* cube_ket,
        const double& max_distance,
        std::list<ReducedTerm>& neighbors) const;
};
} // namespace tbplas::builder
#endif