from math import sqrt
from random import randint
import matplotlib.pyplot as plt

MAX_POINTS = 3 # nbre de points max. par feuille

class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def distance(self, autre_point):
        return sqrt((self.x - autre_point.x)**2 + (self.y - autre_point.y)**2)

class QuadTree:

    def __init__(self, x0, y0, l, h):
        self.x0 = x0
        self.y0 = y0
        self.l = l
        self.h = h
        self.NO = None
        self.NE = None
        self.SE = None
        self.SO = None
        self.points = []

    def is_in(self, x0, y0, l, h, point):
        return point.x >= x0 and point.x <= (x0 + l) and point.y >= y0 and point.y <= (y0 + h)

    def inserer(self, points):
        if len(points) <= MAX_POINTS: # si on a encore de la place dans le quadtree,
            self.points = points  # alors, on stocke les points

        else: # sinon, on le divise :

            # 1. répartition des points dans la liste de leur région respective
            points_NO = []
            points_NE = []
            points_SE = []
            points_SO = []

            for point in points:
                # SO
                if self.is_in(self.x0, self.y0, self.l/2, self.h/2, point):
                    points_SO.append(point)
                # SE
                elif self.is_in(self.x0+self.l/2, self.y0, self.l/2, self.h/2, point):
                    points_SE.append(point)
                # NE
                elif self.is_in(self.x0+self.l/2, self.y0+self.h/2, self.l/2, self.h/2, point):
                    points_NE.append(point)
                # NO
                else:
                    points_NO.append(point)


            # 2. création des quadtree-fils, et insertion de leur liste de points respectives
            self.SO = QuadTree(self.x0, self.y0, self.l/2, self.h/2)
            self.SO.inserer(points_SO)

            self.SE = QuadTree(self.x0+self.l/2, self.y0, self.l/2, self.h/2)
            self.SE.inserer(points_SE)

            self.NE = QuadTree(self.x0+self.l/2, self.y0+self.h/2, self.l/2, self.h/2)
            self.NE.inserer(points_NE)

            self.NO = QuadTree(self.x0, self.y0+self.h/2, self.l/2, self.h/2)
            self.NO.inserer(points_NO)

    def trouver_voisins(self, point):
        if self.SO is None:
            return self.points
        else:
            if self.is_in(self.x0, self.y0, self.l/2, self.h/2, point):
                return self.SO.trouver_voisins(point)
            elif self.is_in(self.x0+self.l/2, self.y0, self.l/2, self.h/2, point):
                return self.SE.trouver_voisins(point)
            elif self.is_in(self.x0+self.l/2, self.y0+self.h/2, self.l/2, self.h/2, point):
                return self.NE.trouver_voisins(point)
            else:
                return self.NO.trouver_voisins(point)

    def _afficher(self):
        plt.plot([self.x0, self.x0+self.l, self.x0+self.l,self.x0, self.x0],[self.y0,self.y0,self.y0+self.h,self.y0+self.h, self.y0], color="black")
        if self.NO is not None:
            self.SO._afficher()
            self.NO._afficher()
            self.SE._afficher()
            self.NE._afficher()
        else:
            x = []
            y = []
            for p in self.points:
                x.append(p.x)
                y.append(p.y)
            plt.scatter(x, y, color='b')

    def afficher(self):
        self._afficher()
        plt.show()

    def afficher_voisins(self, point):
        self._afficher()

        plt.scatter([point.x], [point.y], color="g")

        x = []
        y = []
        for p in self.trouver_voisins(point):
            x.append(p.x)
            y.append(p.y)
        plt.scatter(x, y, color='r')

        plt.show()






points = [Point(randint(1,254), randint(1,254)) for i in range(100)]

qt = QuadTree(0, 0, 255, 255)
qt.inserer(points)
#qt.afficher()

point = Point(randint(1,254), randint(1,254))

qt.afficher_voisins(point)