import numpy as np
from python_prtree import PRTree2D

# Create rectangles: [xmin, ymin, xmax, ymax]
rects = np.array([
    [0.0, 0.0, 1.0, 0.5],  # Rectangle 1
    [1.0, 1.5, 1.2, 3.0],  # Rectangle 2
])
indices = np.array([1, 2])

# Build the tree
tree = PRTree2D(indices, rects)

# Query: find rectangles overlapping with [0.5, 0.2, 0.6, 0.3]
result = tree.query([0.5, 0.2, 0.6, 0.3])
print(result)  # [1]

# Batch query (faster for multiple queries)
queries = np.array([
    [0.5, 0.2, 0.6, 0.3],
    [0.8, 0.5, 1.5, 3.5],
])
results = tree.batch_query(queries)
print(results)  # [[1], [1, 2]]

from python_prtree import PRTree2D, PRTree3D, PRTree4D

tree2d = PRTree2D(indices, boxes_2d)  # [xmin, ymin, xmax, ymax]
tree3d = PRTree3D(indices, boxes_3d)  # [xmin, ymin, zmin, xmax, ymax, zmax]
tree4d = PRTree4D(indices, boxes_4d)  # 4D boxes

# Query with point coordinates
result = tree.query([0.5, 0.5])        # Returns indices
result = tree.query(0.5, 0.5)          # Varargs also supported (2D only)

# Insert new rectangle
tree.insert(3, np.array([1.0, 1.0, 2.0, 2.0]))

# Remove rectangle by index
tree.erase(2)

# Rebuild for optimal performance after many updates
tree.rebuild()

# Store any picklable Python object with rectangles
tree = PRTree2D()
tree.insert(bb=[0, 0, 1, 1], obj={"name": "Building A", "height": 100})
tree.insert(bb=[2, 2, 3, 3], obj={"name": "Building B", "height": 200})

# Query and retrieve objects
results = tree.query([0.5, 0.5, 2.5, 2.5], return_obj=True)
print(results)  # [{'name': 'Building A', 'height': 100}, {'name': 'Building B', 'height': 200}]

# Find all pairs of intersecting rectangles
pairs = tree.query_intersections()
print(pairs)  # numpy array of shape (n_pairs, 2)
# [[1, 3], [2, 5], ...]  # pairs of indices that intersect

# Save tree to file
tree.save('spatial_index.bin')

# Load from file
tree = PRTree2D('spatial_index.bin')

# Or load later
tree = PRTree2D()
tree.load('spatial_index.bin')

# Correct
tree.insert(1, [0, 0, 1, 1])  # xmin=0 < xmax=1, ymin=0 < ymax=1

# Wrong - will raise error
tree.insert(1, [1, 1, 0, 0])  # xmin > xmax, ymin > ymax

tree = PRTree2D()
result = tree.query([0, 0, 1, 1])  # Returns []
results = tree.batch_query(queries)  # Returns [[], [], ...]

PRTree2D(indices=None, boxes=None)
PRTree2D(filename)  # Load from file