import matplotlib.pyplot as mplot
import math

def meanNearestNeighbors(x, y, data_x, data_y, pop_change, k):
    sortedDistances = {}
    for i in range(len(data_x)):
        distance = computeDistance(x,y, data_x[i], data_y[i])
        sortedDistances[distance] = pop_change[i]

    sortedKeys = sorted(sortedDistances.keys())
    sum = 0
    for i in range(k):
        sum += sortedDistances[sortedKeys[i]]

    return sum/k

def computeDistance(x1,y1,x2,y2):
    return math.sqrt(math.pow(x1-x2,2)+math.pow(y1-y2,2))

print("Please input a k-value:")
k = int(input())

us_outline_file = open("us_outline.csv", "r");
us_x = []
us_y = []

data_file = open("data.csv", "r");
data_x = []
data_y = []
pop_change = []
lookup = {}

for line in us_outline_file:
    tokens = line.split(",")
    us_x.append(float(tokens[0]))
    us_y.append(float(tokens[1]))

for line in data_file:
    line = line.replace("\n",'')
    if len(line) > 0:
        tokens = line.split(",")
        lookup[tokens[0]] = {}
        lookup[tokens[0]][tokens[1]] = tokens[2]
        data_x.append(float(tokens[0]))
        data_y.append(float(tokens[1]))
        pop_change.append(float(tokens[2]))

recons_x = []
recons_y= []
recons_value = []
for y in range(121):
    for x in range (195):
        recons_x.append(x)
        recons_y.append(y)
        x = str(x)
        y = str(y)
        if x in lookup:
            if y in lookup[x]:
                recons_value.append(lookup[x][y])
            else:
                recons_value.append(meanNearestNeighbors(int(x), int(y), data_x, data_y, pop_change, k))
        else:
                recons_value.append(meanNearestNeighbors(int(x), int(y), data_x, data_y, pop_change, k))

mplot.plot(us_x, us_y,'-ok', marker=None)
mplot.scatter(recons_x, recons_y, c=recons_value, cmap="viridis", marker=".")

print(len(recons_x))
print(len(recons_y))
print(len(recons_value))

mplot.show()