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

MAX_POINTS = 5 # nbre de points max. par feuille

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

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):
        """Teste si un point appartient à une zone donnée
        Entrées :
            x0, y0 = coordonnées du sommet bas-gauche de la zone
            l, h = largeur et hauteur de la zone
        Sortie :
            booléen True si oui, False si non
        """
        return point.x >= x0 and point.x <= (x0 + l) and point.y >= y0 and point.y <= (y0 + h)

    def inserer(self, points):
        """ Insère récursivement une liste de points dans le quadtree.
        Entrée :
            points = liste d'objets Point à insérer
        Sortie :
            aucune
        """

        if .......: # cas de base ( on est arrivés à une zone de taille minimale )
            self.point = points # on stocke les points dans le noeud ( qui est donc une feuille )

        else: # cas général

            # 1. Répartition des points dans la liste de leur sous-zone respective :

            points_NO = []
            points_NE = []
            points_SE = []
            points_SO = []

            for point in points: # pour chaque point de la liste

                if self.is_in(self.x0; self.y0, self.l//2, self.h//2 , point):  # le point est dans la sous-zone SO ?
                    points_SO.append(point)
                elif self.is_in(..................... , point): # le point est dans la sous-zone SE ?
                    points_SE.append(point)
                elif self.is_in(...................... , point): # le point est dans la sous-zone NE ?
                    points_NE.append(point)
                else:  # sinon ( le point est donc dans la sous-zone NO )
                    points_NO.append(point)

            # 2. création des 4 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.NE = QuadTree(.............................)
            ...................

            self.NO = QuadTree(.............................)
            ...................

    def afficher(self):
        """ Affiche le quadtree """
        plt.plot([self.x0, self.x0+self.l, self.x0+self.l,self.x0],[self.y0,self.y0,self.y0+self.h,self.y0+self.h], color="black")
        if self.NO is not None:
            self.SO.afficher()
            self.NO.afficher()
            self.SE.afficher()
            self.NE.afficher()

    def trouver_voisins(self, point):
        """ Recherche récusrivement dans le quadtree les points les plus proches d'un point donné.
        Entrée :
            point = un objet Point représentant le point autour duquel on cherche les voisins
        Sortie :
            la liste des points voisins
        """

        if ................ : # on est arrivé à une feuille,
            return .................. # on renvoie la liste des points stockés dans cette feuille

        else:
            if self.is_in(.................... , point): # le point est dans la sous-zone SO ?
                return ............................... # alors on recherche récursivement dans la sous-zone SO
            elif self.is_in(...................., point): # idem pour les 3 autres zones
                return ......................
            elif self.is_in(...................., point):
                return ......................
            else:
                return .....................

    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()