import random
import matplotlib.pyplot as plt
import numpy as np

def k_voisins(k):
    liste = [] # contiendra les coordonnées et la classe de chaque point
    classes = ['b','r'] # classes possibles pour les points

    fig = plt.figure()

    for i in range(10):
        x = 10*random.random()
        y = 10*random.random()
        liste.append([x,y,classes[0]])
        plt.plot(x,y,'o',color=classes[0])

    for i in range(10):
        x = 10*random.random()
        y = 10*random.random()
        liste.append([x,y,classes[1]])
        plt.plot(x,y,'o',color=classes[1])

    x = 10*random.random()
    y = 10*random.random()
    plt.plot(x,y,'o',color='g')

    distances = [] # chaque item au format [distance,[x,y,classe]]

    for i in range(20):
        d = ((x-liste[i][0])**2+(y-liste[i][1])**2)**0.5
        distances.append([d,liste[i]])

    distances.sort(key = lambda col: col[0]) # on trie suivant la première colonne de la liste


    count = [0,0]
    for i in range(k):
        plt.plot([x,distances[i][1][0]],[y,distances[i][1][1]],distances[i][1][2]+'--',lw=1)
        if distances[i][1][2] == classes[0]:
            count[0] += 1
        else:
            count[1] +=1

    if count[0] > count[1]:
        classe = classes[0]
    if count[1] > count[0]:
        classe = classes[1]
    if count[0] == count[1]:
        classe = classes[0]+' et '+classes[1]

    plt.title("Pour k = "+str(k)+", la classe du point vert est : "+classe)
    #plt.show()
    fig.savefig('k-voisins-'+str(k)+'.png', dpi=fig.dpi)

for k  in [3,4,5,6,7]:
    k_voisins(k)