import numpy as np
from scipy.spatial import KDTree

np.set_printoptions(precision=3, suppress=True)

points = np.array([
    [0.0, 0.0],
    [2.0, 3.0],
    [2.0, 4.0],
    [5.0, 4.0],
    [8.0, 2.0],
])
labels = np.array(["depot", "sensor-a", "sensor-b", "sensor-c", "sensor-d"])

queries = np.array([
    [2.2, 3.9],
    [7.4, 2.2],
])

tree = KDTree(points)
distances, indexes = tree.query(queries, k=2)

for query_number, query in enumerate(queries, start=1):
    print(f"query {query_number}: {query}")
    for rank, (distance, index) in enumerate(
        zip(distances[query_number - 1], indexes[query_number - 1]),
        start=1,
    ):
        print(
            f"  {rank}: {labels[index]} index={index} "
            f"point={points[index]} distance={distance:.3f}"
        )

manual_distance = np.linalg.norm(points[indexes[0, 0]] - queries[0])
print(f"\nmanual check for query 1 nearest: {manual_distance:.3f}")

bounded_distances, bounded_indexes = tree.query(
    queries[0],
    k=3,
    distance_upper_bound=1.0,
)
print("\nwithin distance 1.0 from query 1")
print("distances:", bounded_distances)
print("indexes:", bounded_indexes)