#include "kdtree.h"
#include "base/consts.h"
namespace tbplas::builder {

Term::Term(const rn_t& rn, const pair_t& pair, const Eigen::Vector3d& rij, const double& distance)
    : rn(rn)
    , pair(pair)
    , rij(rij)
    , distance(distance)
{
}

bool Term::operator<(const Term& other) const
{
    return distance < other.distance;
}

bool Term::operator>(const Term& other) const
{
    return distance > other.distance;
}

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,
    cell_index_t b_max,
    cell_index_t c_max,
    double max_distance)
{
    // Generate rn similar to Python code
    cell_index_t num_rn = (2 * a_max + 1) * (2 * b_max + 1) * (2 * c_max + 1);
    std::vector<rn_t> neighbor_rn(num_rn);
    cell_index_t counter = 0;
    for (cell_index_t ia = -a_max; ia <= a_max; ++ia) {
        for (cell_index_t ib = -b_max; ib <= b_max; ++ib) {
            for (cell_index_t ic = -c_max; ic <= c_max; ++ic) {
                neighbor_rn[counter] = std::make_tuple(ia, ib, ic);
                ++counter;
            }
        }
    }

    // Declare search parameters
    nanoflann::SearchParameters params;
    std::list<Term> neighbors;

    // Loop over all rn
    for (const auto& rn : neighbor_rn) {
        Eigen::Vector3d rn_frac = rn2vec(rn);
        Eigen::Vector3d rn_cart = (rn_frac.transpose() * lat_ket_nm).transpose();
        Eigen::Matrix3Xd pos_ket_rn = pos_ket_nm.colwise() + rn_cart;

        // Loop over all query points (assuming 3D)
        for (size_t orb_j = 0; orb_j < pos_ket_rn.cols(); ++orb_j) {
            std::vector<double> query_pt = { pos_ket_rn(0, orb_j), pos_ket_rn(1, orb_j), pos_ket_rn(2, orb_j) };

            // Declare and initialize results container
            std::vector<nanoflann::ResultItem<unsigned int, double>> ret_matches;

            // Perform radius search
            tree_bra.radiusSearch(&query_pt[0], max_distance * max_distance, ret_matches, params);

            // Process results, excluding the query point itself (distance = 0)
            for (const auto& match : ret_matches) {
                if (match.second > base::INFINITESIMAL) {
                    size_t orb_i = match.first;
                    double distance = std::sqrt(match.second);
                    Eigen::Vector3d rij = pos_ket_rn.col(orb_j) - pos_bra_nm.col(orb_i);
                    neighbors.emplace_back(rn, std::make_pair(orb_i, orb_j), rij, distance);
                }
            }
        }
    }

    // Sort the neighbors by distance
    neighbors.sort();
    return neighbors;
}
} // namespace tbplas::builder