"""
Térinformatikai algoritmusok
Térbeli indexelés
Máté Cserép <mcserep@inf.elte.hu>
"""

import shapefile
import scipy.spatial # pip3 install scipy
import pyqtree       # pip3 install pyqtree

# Open a shapefile (.shp)
# Source: http://gis.inf.elte.hu/files/public/elte-telepulesek
sf = shapefile.Reader('TELEPULESEK.shp')

# Fetch points for cities
points = [shape.points[0] for shape in sf.shapes()]
# Calculating the minimal bounding box
min_x = min(points, key = lambda p: p[0])[0]
max_x = max(points, key = lambda p: p[0])[0]
min_y = min(points, key = lambda p: p[1])[1]
max_y = max(points, key = lambda p: p[1])[1]

print("Number of points: {0}".format(len(points)))
print("Bounding box: ({0}, {1}) - ({2}, {3})".format(min_x, min_y, max_x, max_y))

print("\n --- KdTree ---")
# Create a point which we will query later
#query_point = points[10] # BAD! No cloning, but reference copying
query_point = list(points[10]) # clone the point!
query_point[0] += 1
query_point[1] += 2
print("Query point: {0}".format(query_point))

# Construct KD-Tree
kdtree = scipy.spatial.KDTree(points)

# Query the closest neighbor to the query point
dist, idx = kdtree.query(query_point)
print("Closest neighbor: distance = {0}, index = {1}, point = {2}".format(dist, idx, points[idx]))

# Query the 3 closest neighbors to the query point
dist, idx = kdtree.query(query_point, k = 3)
print("3 closest neighbors:")
for i in range(0, len(dist)):
    print("{0}. neighbor: distance = {1}, index = {2}, point = {3}".format(i+1, dist[i], idx[i], points[idx[i]]))
    
# Query the 100 closest neighbors to the query point within 10km
dist, idx = kdtree.query(query_point, k = 100, distance_upper_bound = 10000)
print("100 closest neighbors within 10km:")
for i in range(0, len(dist)):
    # We will only find 13 neighbors in a 10km range, but the idx list has 100 elements.
    # For the invalid 87 elements the idx[i] is not a valid index, but instead equals to len(points).
    # So with a simple check we can detect the end of the valid results.
    if idx[i] == len(points):
        break
        
    print("{0}. neighbor: distance = {1}, index = {2}, point = {3}".format(i+1, dist[i], idx[i], points[idx[i]]))


print("\n ---QuadTree ---")
# Create a 10x10km query area around a point
query_area_size = 10000
query_area = (
    points[10][0] - query_area_size/2,
    points[10][1] - query_area_size/2,
    points[10][0] + query_area_size/2,
    points[10][1] + query_area_size/2
)
print("Query area: {0}, size = {1} km".format(query_area, query_area_size / 1000))

# Construct the Quad-tree
quadtree = pyqtree.Index(bbox=(min_x, min_y, max_x, max_y))
for point in points:
    quadtree.insert(point, point) # object, bbox
    # For a polygon, the first argument should be the polygon,
    # and the second argument should be the bounding box of the polygon.

# Query the points in the area
matches = quadtree.intersect(query_area)
print("Matches: {0}".format(matches))