import math
import matplotlib.pyplot as mplot

def kNearestNeighbor(dataPoints, x, y, k):
    distances = {}

    # go through each data point and compute distances
    for key in dataPoints.keys():
        dataPoint = dataPoints[key]
        distance = math.sqrt(math.pow(x - float(dataPoint[0]),2) + math.pow(y - float(dataPoint[1]),2))
        distances[distance] = float(dataPoint[2])

    total = 0

    for key in sorted(distances.keys())[0:k]:
        total += distances[key]

    return total//k


print("Input a k value:")
kVal = input()

kVal = int(kVal)

usOutline = open("us_outline.csv", "r");
data = open("data.csv", "r");

usX = []
usY = []
dataX = []
dataY = []
migration = []

migrationTable = {}

# draw US outline

for line in usOutline:
    coordinates = line.split(",")

    xCoord = float(coordinates[0])
    yCoord = float(coordinates[1])

    usX.append(xCoord)
    usY.append(yCoord)

# create a hash table to look up information for each data point.

for line in data:
    line = line.replace("\n",'')

    if len(line) != 0:
        locationInfo = line.split(",")
        migrationTable[locationInfo[0] + "," + locationInfo[1]] = locationInfo

# perform reconstruction

resultX = []
resultY= []
resultValue = []

for y in range(121):
    for x in range (195):
        resultX.append(x)
        resultY.append(y)
        point = str(x) + "," + str(y)

        if point in migrationTable:
            resultValue.append(migrationTable[point][2])
        else:
            resultValue.append(kNearestNeighbor(migrationTable, x, y, kVal))

mplot.plot(usX, usY,'-ok', marker=".")

mplot.scatter(resultX, resultY, c=resultValue, cmap="viridis", marker=".")

mplot.show()