1"""
2strtree
3=======
4
5Index geometry objects for efficient lookup of nearby or
6nearest neighbors. Home of the `STRtree` class which is
7an interface to the query-only GEOS R-tree packed using
8the Sort-Tile-Recursive algorithm [1]_.
9
10.. autoclass:: STRtree
11    :members:
12
13References
14----------
15  .. [1]  Leutenegger, Scott & Lopez, Mario & Edgington, Jeffrey. (1997).
16     "STR: A Simple and Efficient Algorithm for R-Tree Packing." Proc.
17     VLDB Conf. 497-506. 10.1109/ICDE.1997.582015.
18     https://www.cs.odu.edu/~mln/ltrs-pdfs/icase-1997-14.pdf
19"""
20
21import ctypes
22import logging
23from typing import Any, ItemsView, Iterable, Iterator, Optional, Sequence, Tuple, Union
24import sys
25
26from shapely.geometry.base import BaseGeometry
27from shapely.geos import lgeos
28
29log = logging.getLogger(__name__)
30
31
32class STRtree:
33    """An STR-packed R-tree spatial index.
34
35    An index is initialized from a sequence of geometry objects and
36    optionally an sequence of items. The items, if provided, are stored
37    in nodes of the tree. If items are not provided, the indices of the
38    geometry sequence will be used instead.
39
40    Stored items and corresponding geometry objects can be spatially
41    queried using another geometric object.
42
43    The tree is immutable and query-only, meaning that once created
44    nodes cannot be added or removed.
45
46    Parameters
47    ----------
48    geoms : sequence
49        A sequence of geometry objects.
50    items : sequence, optional
51        A sequence of objects which typically serve as identifiers in an
52        application. This sequence must have the same length as geoms.
53
54    Attributes
55    ----------
56    node_capacity : int
57        The maximum number of items per node. Default: 10.
58
59    Examples
60    --------
61    Creating an index of polygons:
62
63    >>> from shapely.strtree import STRtree
64    >>> from shapely.geometry import Polygon
65    >>>
66    >>> polys = [Polygon(((0, 0), (1, 0), (1, 1))),
67    ...          Polygon(((0, 1), (0, 0), (1, 0))),
68    ...          Polygon(((100, 100), (101, 100), (101, 101)))]
69    >>> tree = STRtree(polys)
70    >>> query_geom = Polygon(((-1, -1), (2, 0), (2, 2), (-1, 2)))
71    >>> result = tree.query(query_geom)
72    >>> polys[0] in result
73    True
74    >>> polys[1] in result
75    True
76    >>> polys[2] in result
77    False
78
79    Notes
80    -----
81    The class maintains a reverse mapping of items to geometries. These
82    items must therefore be hashable. The tree is filled using the
83    Sort-Tile-Recursive [1]_ algorithm.
84
85    References
86    ----------
87    .. [1] Leutenegger, Scott T.; Edgington, Jeffrey M.; Lopez, Mario A.
88       (February 1997). "STR: A Simple and Efficient Algorithm for
89       R-Tree Packing".
90       https://ia600900.us.archive.org/27/items/nasa_techdoc_19970016975/19970016975.pdf
91
92    """
93
94    def __init__(
95        self,
96        geoms: Iterable[BaseGeometry],
97        items: Iterable[Any] = None,
98        node_capacity: int = 10,
99    ):
100        self.node_capacity = node_capacity
101
102        # Keep references to geoms
103        self._geoms = list(geoms)
104        # Default enumeration index to store in the tree
105        self._idxs = list(range(len(self._geoms)))
106
107        # handle items
108        self._has_custom_items = items is not None
109        if not self._has_custom_items:
110            items = self._idxs
111        self._items = items
112
113        # initialize GEOS STRtree
114        self._tree = lgeos.GEOSSTRtree_create(self.node_capacity)
115        i = 0
116        for idx, geom in zip(self._idxs, self._geoms):
117            # filter empty geometries out of the input
118            if geom is not None and not geom.is_empty:
119                lgeos.GEOSSTRtree_insert(self._tree, geom._geom, ctypes.py_object(idx))
120                i += 1
121        self._n_geoms = i
122
123    def __reduce__(self):
124        if self._has_custom_items:
125            return STRtree, (self._geoms, self._items)
126        else:
127            return STRtree, (self._geoms, )
128
129    def __del__(self):
130        if self._tree is not None:
131            try:
132                lgeos.GEOSSTRtree_destroy(self._tree)
133            except AttributeError:
134                pass  # lgeos might be empty on shutdown.
135
136            self._tree = None
137
138    def _query(self, geom):
139        if self._n_geoms == 0:
140            return []
141
142        result = []
143
144        def callback(item, userdata):
145            idx = ctypes.cast(item, ctypes.py_object).value
146            result.append(idx)
147
148        lgeos.GEOSSTRtree_query(self._tree, geom._geom, lgeos.GEOSQueryCallback(callback), None)
149        return result
150
151    def query_items(self, geom: BaseGeometry) -> Sequence[Any]:
152        """Query for nodes which intersect the geom's envelope to get
153        stored items.
154
155        Items are integers serving as identifiers for an application.
156
157        Parameters
158        ----------
159        geom : geometry object
160            The query geometry.
161
162        Returns
163        -------
164        An array or list of items stored in the tree.
165
166        Note
167        ----
168        A geometry object's "envelope" is its minimum xy bounding
169        rectangle.
170
171        Examples
172        --------
173        A buffer around a point can be used to control the extent
174        of the query.
175
176        >>> from shapely.strtree import STRtree
177        >>> from shapely.geometry import Point
178        >>> points = [Point(i, i) for i in range(10)]
179        >>> tree = STRtree(points)
180        >>> query_geom = Point(2,2).buffer(0.99)
181        >>> [o.wkt for o in tree.query(query_geom)]
182        ['POINT (2 2)']
183        >>> query_geom = Point(2, 2).buffer(1.0)
184        >>> [o.wkt for o in tree.query(query_geom)]
185        ['POINT (1 1)', 'POINT (2 2)', 'POINT (3 3)']
186
187        A subsequent search through the returned subset using the
188        desired binary predicate (eg. intersects, crosses, contains,
189        overlaps) may be necessary to further filter the results
190        according to their specific spatial relationships.
191
192        >>> [o.wkt for o in tree.query(query_geom) if o.intersects(query_geom)]
193        ['POINT (2 2)']
194
195        """
196        result = self._query(geom)
197        if self._has_custom_items:
198            return [self._items[i] for i in result]
199        else:
200            return result
201
202    def query_geoms(self, geom: BaseGeometry) -> Sequence[BaseGeometry]:
203        """Query for nodes which intersect the geom's envelope to get
204        geometries corresponding to the items stored in the nodes.
205
206        Parameters
207        ----------
208        geom : geometry object
209            The query geometry.
210
211        Returns
212        -------
213        An array or list of geometry objects.
214
215        """
216        result = self._query(geom)
217        return [self._geoms[i] for i in result]
218
219    def query(self, geom: BaseGeometry) -> Sequence[BaseGeometry]:
220        """Query for nodes which intersect the geom's envelope to get
221        geometries corresponding to the items stored in the nodes.
222
223        This method is an alias for query_geoms. It may be removed in
224        version 2.0.
225
226        Parameters
227        ----------
228        geom : geometry object
229            The query geometry.
230
231        Returns
232        -------
233        An array or list of geometry objects.
234
235        """
236        return self.query_geoms(geom)
237
238    def _nearest(self, geom, exclusive):
239        envelope = geom.envelope
240
241        def callback(item1, item2, distance, userdata):
242            try:
243                callback_userdata = ctypes.cast(userdata, ctypes.py_object).value
244                idx = ctypes.cast(item1, ctypes.py_object).value
245                geom2 = ctypes.cast(item2, ctypes.py_object).value
246                dist = ctypes.cast(distance, ctypes.POINTER(ctypes.c_double))
247                if callback_userdata["exclusive"] and self._geoms[idx].equals(geom2):
248                    dist[0] = sys.float_info.max
249                else:
250                    lgeos.GEOSDistance(self._geoms[idx]._geom, geom2._geom, dist)
251
252                return 1
253            except Exception:
254                log.exception("Caught exception")
255                return 0
256
257        item = lgeos.GEOSSTRtree_nearest_generic(
258            self._tree,
259            ctypes.py_object(geom),
260            envelope._geom,
261            lgeos.GEOSDistanceCallback(callback),
262            ctypes.py_object({"exclusive": exclusive}),
263        )
264        return ctypes.cast(item, ctypes.py_object).value
265
266    def nearest_item(
267        self, geom: BaseGeometry, exclusive: bool = False
268    ) -> Union[Any, None]:
269        """Query the tree for the node nearest to geom and get the item
270        stored in the node.
271
272        Items are integers serving as identifiers for an application.
273
274        Parameters
275        ----------
276        geom : geometry object
277            The query geometry.
278        exclusive : bool, optional
279            Whether to exclude the item corresponding to the given geom
280            from results or not.  Default: False.
281
282        Returns
283        -------
284        Stored item or None.
285
286        None is returned if this index is empty. This may change in
287        version 2.0.
288
289        Examples
290        --------
291        >>> from shapely.strtree import STRtree
292        >>> from shapely.geometry import Point
293        >>> tree = STRtree([Point(i, i) for i in range(10)])
294        >>> tree.nearest(Point(2.2, 2.2)).wkt
295        'POINT (2 2)'
296
297        Will only return one object:
298
299        >>> tree = STRtree ([Point(0, 0), Point(0, 0)])
300        >>> tree.nearest(Point(0, 0)).wkt
301        'POINT (0 0)'
302
303        """
304        if self._n_geoms == 0:
305            return None
306
307        result = self._nearest(geom, exclusive)
308        if self._has_custom_items:
309            return self._items[result]
310        else:
311            return result
312
313    def nearest_geom(
314        self, geom: BaseGeometry, exclusive: bool = False
315    ) -> Union[BaseGeometry, None]:
316        """Query the tree for the node nearest to geom and get the
317        geometry corresponding to the item stored in the node.
318
319        Parameters
320        ----------
321        geom : geometry object
322            The query geometry.
323        exclusive : bool, optional
324            Whether to exclude the given geom from results or not.
325            Default: False.
326
327        Returns
328        -------
329        BaseGeometry or None.
330
331        None is returned if this index is empty. This may change in
332        version 2.0.
333
334        """
335        result = self._nearest(geom, exclusive)
336        return self._geoms[result]
337
338    def nearest(
339        self, geom: BaseGeometry, exclusive: bool = False
340    ) -> Union[BaseGeometry, None]:
341        """Query the tree for the node nearest to geom and get the
342        geometry corresponding to the item stored in the node.
343
344        This method is an alias for nearest_geom. It may be removed in
345        version 2.0.
346
347        Parameters
348        ----------
349        geom : geometry object
350            The query geometry.
351        exclusive : bool, optional
352            Whether to exclude the given geom from results or not.
353            Default: False.
354
355        Returns
356        -------
357        BaseGeometry or None.
358
359        None is returned if this index is empty. This may change in
360        version 2.0.
361
362        """
363        return self.nearest_geom(geom, exclusive=exclusive)
364