From 006c8c97fa05cb89b9a6af7d64e674728888e6e8 Mon Sep 17 00:00:00 2001
From: Pedro Maciel <Pedro.Maciel@ecmwf.int>
Date: Tue, 15 Aug 2023 09:54:55 +0100
Subject: [PATCH] eckit::geometry

---
 src/eckit/geometry/Grid.cc  | 24 +++++++++++++++++++
 src/eckit/geometry/Grid.h   |  4 ++++
 src/eckit/geometry/Search.h | 48 +++++++++++++++++++++++--------------
 3 files changed, 58 insertions(+), 18 deletions(-)

diff --git a/src/eckit/geometry/Grid.cc b/src/eckit/geometry/Grid.cc
index 9571c574a..560f026d6 100644
--- a/src/eckit/geometry/Grid.cc
+++ b/src/eckit/geometry/Grid.cc
@@ -71,6 +71,30 @@ std::vector<Point> Grid::to_points() const {
 }
 
 
+SearchLonLat::Result Grid::nearest(const PointLonLat& p) const {
+    SearchLonLat search;
+
+    size_t index = 0;
+    for (const auto& p : to_points()) {
+        search.insert({std::get<PointLonLat>(p), index});
+    }
+
+    return search.nearestNeighbour(p);
+}
+
+
+SearchLonLat::Results Grid::nearest(const PointLonLat& p, size_t k) const {
+    SearchLonLat search;
+
+    size_t index = 0;
+    for (const auto& p : to_points()) {
+        search.insert({std::get<PointLonLat>(p), index});
+    }
+
+    return search.kNearestNeighbours(p, k);
+}
+
+
 static pthread_once_t __once;
 
 static Mutex* __mutex                                        = nullptr;
diff --git a/src/eckit/geometry/Grid.h b/src/eckit/geometry/Grid.h
index b95ca1c04..ede011a38 100644
--- a/src/eckit/geometry/Grid.h
+++ b/src/eckit/geometry/Grid.h
@@ -24,6 +24,7 @@
 #include "eckit/geometry/Point.h"
 #include "eckit/geometry/Projection.h"
 #include "eckit/geometry/Renumber.h"
+#include "eckit/geometry/Search.h"
 #include "eckit/geometry/area/BoundingBox.h"
 
 
@@ -102,6 +103,9 @@ class Grid {
 
     virtual std::vector<Point> to_points() const;
 
+    virtual SearchLonLat::Result nearest(const PointLonLat&) const;
+    virtual SearchLonLat::Results nearest(const PointLonLat&, size_t k) const;
+
     // -- Overridden methods
     // None
 
diff --git a/src/eckit/geometry/Search.h b/src/eckit/geometry/Search.h
index 31f7cc9af..bec4ca509 100644
--- a/src/eckit/geometry/Search.h
+++ b/src/eckit/geometry/Search.h
@@ -12,11 +12,12 @@
 
 #pragma once
 
+#include <tuple>
+#include <vector>
+
 #include "eckit/container/KDTree.h"
 #include "eckit/container/sptree/SPValue.h"
-#include "eckit/geometry/Point2.h"
-#include "eckit/geometry/Point3.h"
-#include "eckit/geometry/PointLonLat.h"
+#include "eckit/geometry/Point.h"
 #include "eckit/geometry/UnitSphere.h"
 
 
@@ -42,38 +43,49 @@ struct SearchLonLat : Search3 {
     using Point = PointLonLat;
     using Value = SPValue<TT<search::Traits<Point, Payload>, KDMemory>>;
 
+    using Result  = std::tuple<Point, Payload, double>;
+    using Results = std::vector<Result>;
+
     using Search3::Search3;
 
     void insert(const SearchLonLat::Value& value) {
-        Search3::insert({convert(value.point()), value.payload()});
+        Search3::insert({to_cartesian(value.point()), value.payload()});
     }
 
-    template <typename ITER>
-    void build(ITER begin, ITER end) {
-        for (auto it = begin; it != end; ++it) {
-            insert(*it);
+    template <typename Container>
+    void build(const Container& c) {
+        size_t index = 0;
+        for (const auto& p : c) {
+            insert({p, index++});
         }
     }
 
-    template <typename Container>
-    void build(Container& c) {
-        build(c.begin(), c.end());
+    Result nearestNeighbour(const Point& p) {
+        auto n = Search3::nearestNeighbour(to_cartesian(p));
+        return {UnitSphere::convertCartesianToSpherical(n.point()), n.payload(), n.distance()};
     }
 
-    NodeInfo nearestNeighbour(const Point& p) { return Search3::nearestNeighbour(convert(p)); }
-
-    NodeList findInSphere(const Point& p, double radius) {
-        return Search3::findInSphere(convert(p), radius);
+    Results findInSphere(const Point& p, double radius) {
+        return to_spherical(Search3::findInSphere(to_cartesian(p), radius));
     }
 
-    NodeList kNearestNeighbours(const Point& p, size_t k) {
-        return Search3::kNearestNeighbours(convert(p), k);
+    Results kNearestNeighbours(const Point& p, size_t k) {
+        return to_spherical(Search3::kNearestNeighbours(to_cartesian(p), k));
     }
 
 private:
-    static Search3::Point convert(const Point& p) {
+    static Search3::Point to_cartesian(const Point& p) {
         return UnitSphere::convertSphericalToCartesian(p);
     }
+
+    static Results to_spherical(const NodeList& nodes) {
+        Results list;
+        list.reserve(nodes.size());
+        for (const auto& n : nodes) {
+            list.emplace_back(UnitSphere::convertCartesianToSpherical(n.point()), n.payload(), n.distance());
+        }
+        return list;
+    }
 };