"""
Térinformatikai algoritmusok
K-közép klaszterezés
Máté Cserép <mcserep@inf.elte.hu>
"""

import shapefile

import numpy as np
import matplotlib.pyplot as plt

from sklearn.cluster import KMeans # pip3 install scikit-learn

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

# Fetch points for cities (as numpy array)
points = np.array([shape.points[0] for shape in sf.shapes()])

# Plot figure
plt.figure(figsize=(8, 8))

# Cluster the points
pred = KMeans(n_clusters=19).fit_predict(points)
# Put them on the plot
plt.scatter(points[:, 0], points[:, 1], c=pred)

# Display plot
plt.title("Cluster map of the Hungarian cities")
plt.show()