From 33f736ce389795486b027196ec8e8045e24ab4c0 Mon Sep 17 00:00:00 2001
From: Manos Papadakis <papadakm95@gmail.com>
Date: Tue, 14 Nov 2023 22:00:19 +0200
Subject: [PATCH] Update KDTreeArmadilloAdaptor.hpp

---
 .../internal/KDTreeArmadilloAdaptor.hpp       | 76 ++++++++++++++++++-
 1 file changed, 75 insertions(+), 1 deletion(-)

diff --git a/inst/include/internal/KDTreeArmadilloAdaptor.hpp b/inst/include/internal/KDTreeArmadilloAdaptor.hpp
index 06ce5f8..25b210e 100644
--- a/inst/include/internal/KDTreeArmadilloAdaptor.hpp
+++ b/inst/include/internal/KDTreeArmadilloAdaptor.hpp
@@ -50,6 +50,10 @@ namespace Rnanoflann {
         {
             return m_data_matrix.get()(dim,idx);
         }
+
+        colvec col(uword idx) const {
+            return m_data_matrix.get().col(idx);
+        }
         
         template <class BBOX>
         bool kdtree_get_bbox(BBOX& bb) const
@@ -60,7 +64,7 @@ namespace Rnanoflann {
 
 
     // Define an Armadillo KDTreeAdaptor class
-    template <class MatrixType, class Distance, bool Square = false, int DIM = -1>
+    template <class MatrixType, class Distance, bool Square, int DIM = -1>
     struct KDTreeArmadilloAdaptor2
     {
         using self_t = KDTreeArmadilloAdaptor2<MatrixType, Distance, Square, DIM>;
@@ -96,6 +100,10 @@ namespace Rnanoflann {
         {
             return m_data_matrix.get().n_cols;
         }
+
+        colvec col(uword idx) const {
+            return m_data_matrix.get().col(idx);
+        }
         
         num_t kdtree_get_pt(uword idx, size_t dim) const
         {
@@ -108,4 +116,70 @@ namespace Rnanoflann {
             return false; // Optional bounding-box computation (not used in this example)
         }
     };
+
+    template <class MatrixType, class Distance, int DIM = -1>
+    struct KDTreeArmadilloAdaptor3
+    {
+        using self_t = KDTreeArmadilloAdaptor3<MatrixType, Distance, DIM>;
+        using num_t = typename MatrixType::elem_type;
+        using IndexType =  uword;
+        
+        using metric_t = typename Distance::template traits<num_t, self_t>::distance_t;// You can change the distance metric as needed.
+        
+        using index_t = nanoflann::KDTreeSingleIndexAdaptor<metric_t, self_t, DIM, IndexType>;
+        
+        index_t* index_;
+
+        const double p;
+        const double p_1;
+        
+        const std::reference_wrapper<const MatrixType> m_data_matrix;
+        
+        explicit KDTreeArmadilloAdaptor3(const uword dimensionality, const std::reference_wrapper<const MatrixType>& mat, const double p,const uword leafs = 10)
+            : p(p), p_1(1.0/p), m_data_matrix(mat)
+        {
+            const auto dims = mat.get().n_rows; // Assumes column-major Armadillo matrix
+            if (static_cast<uword>(dims) != dimensionality)
+                throw std::runtime_error("Error: 'dimensionality' must match the column count in the data matrix");
+            if (DIM > 0 && static_cast<int>(dims) != DIM)
+                throw std::runtime_error("Data set dimensionality does not match the 'DIM' template argument");
+            index_ = new index_t(dims, *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leafs));
+            index_->buildIndex();
+        }
+        
+        ~KDTreeArmadilloAdaptor3() { delete index_; }
+        
+        const self_t& derived() const { return *this; }
+        self_t& derived() { return *this; }
+        
+        uword kdtree_get_point_count() const
+        {
+            return m_data_matrix.get().n_cols;
+        }
+        
+        num_t kdtree_get_pt(uword idx, size_t dim) const
+        {
+            return m_data_matrix.get()(dim,idx);
+        }
+
+        colvec col(uword idx) const {
+            return m_data_matrix.get().col(idx);
+        }
+
+        const double getP() const {
+            return p;
+        }
+
+        const double getP_1() const {
+            return p_1;
+        }
+        
+        template <class BBOX>
+        bool kdtree_get_bbox(BBOX& bb) const
+        {
+            return false; // Optional bounding-box computation (not used in this example)
+        }
+    };
+
+
 };
\ No newline at end of file