"""
Térinformatikai algoritmusok
Gráfalgoritmusok II.: minimális költségű feszítőfák
Máté Cserép <mcserep@inf.elte.hu>
"""

import shapefile                # pip3 install pyshp
import networkx as nx           # pip3 install networkx
import matplotlib.pyplot as plt	# pip3 install matplotlib

# Open a shapefile (.shp)
# dBase file of attributes (.dbf) is automatically detected by name convention
# Source: http://gis.inf.elte.hu/files/public/elte-telepulesek
sf = shapefile.Reader('TELEPULESEK.shp')

# Check whether files contains points
if sf.shapeType == shapefile.POINT:
	print("This file contains points")
else:
    print("This file does not contain points")
    quit()
    
# Print some file level metadata
print("Number of cities: {0}".format(len(sf.shapes())))
print("Attributes: {0}".format(sf.fields))

print("------------------")

# Create empty, undirected graph
graph = nx.Graph()

# Iterate through all vector features (cities)
print("Reading shapefile ...")
for sr in sf.shapeRecords():
	shape  = sr.shape  # geometry
	record = sr.record # attributes
	
    # We could see from the attribute list that the name is the 3rd attribute
	name = record[2] 
	if isinstance(name, bytes):
		name = str(name, "iso-8859-2")
	
    # Add the city to the graph with its name
    # Store its location as addditional data
	graph.add_node(name, location = shape.points[0])
	
# Create a complete graph (add all possible edges)
print("Construting complete graph ...")
for city_from in graph.nodes:
	location_from = graph.nodes[city_from]['location']
	for city_to in graph.nodes:
		location_to = graph.nodes[city_to]['location']
		if city_from < city_to: # we do not need to add all edges twice
            # Add edge to the graph with distance as its cost
			graph.add_edge(city_from, city_to, 
				distance = pow(location_from[0] - location_to[0], 2) + pow(location_from[1] - location_to[1], 2))

# Calculates the minimum spanning tree as a new graph
print("Calculating minimum spanning tree ...")
spanning_tree = nx.minimum_spanning_tree(graph, weight = 'distance')

# Start new plot figure
print("Drawing plot ...")
plt.figure()

# Plot all edges as black lines in the MST
for edge in spanning_tree.edges:
	city_from = edge[0]
	city_to   = edge[1]
	
	location_from = spanning_tree.nodes[city_from]['location']
	location_to   = spanning_tree.nodes[city_to]['location']
	plt.plot([location_from[0], location_to[0]], [location_from[1], location_to[1]], color = 'black')

# Plot all cities as red dots
for city in spanning_tree.nodes:
	location = spanning_tree.nodes[city]['location']
	plt.plot(location[0], location[1], color = 'red', marker = 'o', markersize = 1)
	
# Display plot
print("Ready.")
plt.show()