From f79b4cf81740e4f9f72dd463e5475a8570c3fb4a Mon Sep 17 00:00:00 2001
From: lebuller <lancebullerwell12@gmail.com>
Date: Thu, 9 Nov 2023 11:58:14 -0500
Subject: [PATCH] Implements getNeighbor algorithm for linkedCellLists (needs
 testing)

---
 core/src/Cabana_LinkedCellList.hpp | 37 +++++++++++++++++++++++-------
 1 file changed, 29 insertions(+), 8 deletions(-)

diff --git a/core/src/Cabana_LinkedCellList.hpp b/core/src/Cabana_LinkedCellList.hpp
index 9c2e6e775..8603e3fd0 100644
--- a/core/src/Cabana_LinkedCellList.hpp
+++ b/core/src/Cabana_LinkedCellList.hpp
@@ -696,15 +696,36 @@ class NeighborList<LinkedCellList<DeviceType, Scalar>>
 
     //! Get the id for a neighbor for a given particle index and the index of
     //! the neighbor relative to the particle.
-    KOKKOS_INLINE_FUNCTION
-    static std::size_t getNeighbor( const list_type& list, const std::size_t,
-                                    const std::size_t neighbor_index,
-                                    const bool sorted = true )
+    template <class CellIndexType>
+    KOKKOS_INLINE_FUNCTION static std::size_t
+    getNeighbor( const list_type& list, const CellIndexType cell,
+                 const std::size_t particle_index,
+                 const std::size_t neighbor_index, const bool sorted = true )
     {
-        if ( sorted )
-            return neighbor_index;
-        else
-            return list.permutation( neighbor_index );
+        int total_count = 0;
+        int previous_count = 0;
+        int imin, imax, jmin, jmax, kmin, kmax;
+        list.getStencilCells( cell( particle_index ), imin, imax, jmin, jmax,
+                              kmin, kmax );
+
+        // Loop over the cell stencil.
+        for ( int i = imin; i < imax; ++i )
+            for ( int j = jmin; j < jmax; ++j )
+                for ( int k = kmin; k < kmax; ++k )
+                {
+                    total_count += list.binSize( i, j, k );
+                    if ( total_count > neighbor_index )
+                    {
+                        int particle_id = list.binOffset( i, j, k ) +
+                                          ( neighbor_index - previous_count );
+                        if ( sorted )
+                            return particle_id;
+                        else
+                            return list.permutation( particle_id );
+                        break;
+                    }
+                    previous_count = total_count;
+                }
     }
 };