/**
 * @brief Interface to the nanoflann KDTree library
 *
 * @author: Zewen Wu (zwwu)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_BUILDER_KDTREE_H
#define TBPLAS_BUILDER_KDTREE_H

#include <list>
#include <vector>

#include <eigen3/Eigen/Dense>
#include <nanoflann/nanoflann.hpp>

#include "base.h"

namespace tbplas::builder {

/// @brief A matrix class template?
template <typename T>
struct VectorOfVectorsAdapter {
    const std::vector<std::vector<T>>& data;

    explicit VectorOfVectorsAdapter(const std::vector<std::vector<T>>& data_)
        : data(data_)
    {
    }

    inline size_t kdtree_get_point_count() const { return data.size(); }

    inline T kdtree_distance(const T* p1, const size_t idx_p2,
        size_t /*size*/) const
    {
        T s = 0;
        for (size_t i = 0; i < data[0].size(); ++i) {
            const T d = p1[i] - data[idx_p2][i];
            s += d * d;
        }
        return s;
    }

    inline T kdtree_get_pt(const size_t idx, int dim) const
    {
        return data[idx][dim];
    }

    template <class BBOX>
    bool kdtree_get_bbox(BBOX& /*bb*/) const
    {
        return false;
    }
};

using kd_tree_t = nanoflann::KDTreeSingleIndexAdaptor<
    nanoflann::L2_Simple_Adaptor<double, VectorOfVectorsAdapter<double>>,
    VectorOfVectorsAdapter<double>, 3>;

/// @brief Datatype for hopping term in results
struct Term {
    rn_t rn = { 0, 0, 0 }; // cell index (a, b, c)
    pair_t pair = { 0, 0 }; // orbital indices
    Eigen::Vector3d rij = Eigen::Vector3d::Zero(); // displacement vector in nm
    double distance = 0.0; // norm of rij in nm

    // Constructor
    Term(const rn_t& rn, const pair_t& pair, const Eigen::Vector3d& rij, const double& distance);

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

/// @brief Find neighbors within cutoff distance
/// @param[in] tree_bra KDTree created from orbital positions of pc_bra
/// @param[in] pos_bra_nm Cartesian coordinates of orbitals of pc_bra in nm in COLUMN-WISE order
/// @param[in] pos_ket_nm Cartesian coordinates of orbitals of pc_ket in nm in COLUMN-WISE order
/// @param[in] lat_ket_nm Cartesian coordinates of lattice vectors of pc_ket in nm in ROW-WISE order
/// @param[in] a_max upper bound of range on a-axis
/// @param[in] b_max upper bound of range on b-axis
/// @param[in] c_max upper bound of range on c-axis
/// @param[in] max_distance  cutoff distance in NM
/// @return list of hopping neighbors
std::list<Term> find_neighbors_within_radius(kd_tree_t& tree_bra,
    const Eigen::Matrix3Xd& pos_bra_nm,
    const Eigen::Matrix3Xd& pos_ket_nm,
    const Eigen::Matrix3d& lat_ket_nm,
    cell_index_t a_max = 0,
    cell_index_t b_max = 0,
    cell_index_t c_max = 0,
    double max_distance = 1.0);
} // namespace tbplas::builder
#endif