import math
import numpy
import random
import pygame
            
class Disque:
    def __init__(self,x,y,vx,vy):
        self.x = x
        self.y = y
        self.vx  = vx
        self.vy = vy
        self.ax = 0.0
        self.ay = 0.0
            
class Cellule:
    def __init__(self):
        self.ensemble_disques = set()
        self.decal_x = 0.0
        self.decal_y = 0.0
    def ajouter(self,indice):
        self.ensemble_disques.add(indice)
    def enlever(self,indice):
        self.ensemble_disques.remove(indice)
             
class Grille:
    def __init__(self,Nc,L):
        self.L = L
        self.Nc = Nc
        self.tableau = []
        for i in range(self.Nc):
            ligne = []
            for j in range(self.Nc):
                ligne.append(Cellule())
            self.tableau.append(ligne)
    def obtenir_cellule(self,i,j):
            decal_x = 0.0
            decal_y = 0.0
            if i<0:
                i += self.Nc
                decal_x = -self.L
            elif i>=self.Nc:
                i -= self.Nc
                decal_x = self.L
            if j<0:
                j += self.Nc
                decal_y = -self.L
            elif j>=self.Nc:
                j -= self.Nc
                decal_y = self.L
            cellule = self.tableau[i][j]
            cellule.decal_x = decal_x
            cellule.decal_y = decal_y
            return cellule
             
class Systeme:
    def __init__(self,Nx,densite):
        self.Nx = Nx
        self.N = Nx*Nx
        self.rayon = 0.5
        self.diam = self.rayon*2
        self.diam2 = self.diam**2
        self.L = math.sqrt(math.pi/densite)*0.5*Nx
        self.aire = self.L*self.L
        self.demi_L = self.L*0.5
        self.Nc = int(self.L/self.diam)
        self.lc = self.L/self.Nc
        self.grille = Grille(self.Nc,self.L)
        self.liste_disques = []
        for d in range(self.N):
            self.liste_disques.append(Disque(0,0,0,0))
        self.energie = 0.0
        self.viriel = 0.0
        self.Ecinetique = 0.0
        self.Epotentielle = 0.0
        self.pression = 0.0

            
    def init_disque(self,indice,x,y,vx,vy):
        disque = self.liste_disques[indice]
        disque.x = x
        disque.y = y
        disque.vx = vx
        disque.vy = vy
        i = int(x/self.lc)
        j = int(y/self.lc)
        cellule = self.grille.obtenir_cellule(i,j)
        cellule.ajouter(indice)

            
    def initialiser(self,vitesse):
        dx = self.L*1.0/(self.Nx)
        print("distance initiale = %f"%dx)
        if dx < self.diam:
            raise Exception("densite trop forte")
        else:
            dy = dx
            x = dx/2
            y = x
            px = 0.0
            py = 0.0
            for k in range(self.N):
                a = random.random()*math.pi*2.0
                vx = vitesse*math.cos(a)
                vy = vitesse*math.sin(a)
                px += vx
                py += vy
                self.init_disque(k,x,y,vx,vy)
                x += dx
                if x > self.L:
                    x = dx/2
                    y += dy
            for k in range(self.N):
                disque = self.liste_disques[k]
                disque.vx -= px/self.N
                disque.vy -= py/self.N 
        self.calculer_forces()
           
    def deplacer_disque(self,disque,indice,x1,y1):
        if x1<0:
            x1 %= self.L
        elif x1>=self.L:
            x1 %= self.L
        if y1<0:
            y1 %= self.L
        elif y1>=self.L:
            y1 %= self.L
        i = int(disque.x/self.lc)
        j = int(disque.y/self.lc)
        i1 = int(x1/self.lc)
        j1 = int(y1/self.lc)
        if i!=i1 or j!=j1:
            cellule = self.grille.obtenir_cellule(i,j)
            cellule.enlever(indice)
            cellule1 = self.grille.obtenir_cellule(i1,j1)
            cellule1.ajouter(indice)
        disque.x = x1
        disque.y = y1

           
    def calculer_forces(self):
        for k in range(self.N):
            disque = self.liste_disques[k]
            disque.ax = 0.0
            disque.ay = 0.0
        self.Epotentielle = 0.0
        self.viriel = 0.0
        for i in range(self.Nc):
            for j in range(self.Nc):
                cellule = self.grille.obtenir_cellule(i,j)
                for k in range(-1,2):
                    for l in range(-1,2):
                        cellule1 = self.grille.obtenir_cellule(i+k,j+l)
                        for indice in cellule.ensemble_disques:
                                for indice1 in cellule1.ensemble_disques:
                                    if indice1<indice:
                                        disque = self.liste_disques[indice]
                                        disque1 = self.liste_disques[indice1]
                                        dx = disque1.x+cellule1.decal_x-disque.x
                                        dy = disque1.y+cellule1.decal_y-disque.y
                                        r2 = dx*dx+dy*dy
                                        if r2 < self.diam2:
                                            ir2 = 1.0/r2
                                            ir6 = ir2*ir2*ir2
                                            f = 12.0*ir6*(ir6-1.0)*ir2
                                            fx = f*dx
                                            fy = f*dy
                                            disque1.ax += fx
                                            disque1.ay += fy
                                            disque.ax -= fx
                                            disque.ay -= fy
                                            self.Epotentielle += ir6*(ir6-2.0)+1.0
                                            self.viriel += 6.0*ir6*(ir6-1.0)

             
    def calculer_cinetique(self):
        self.Ecinetique = 0.0
        for k in range(self.N):
            disque = self.liste_disques[k]
            self.Ecinetique += 0.5*(disque.vx*disque.vx+disque.vy*disque.vy)
        self.pression = (self.Ecinetique+self.viriel)/self.aire
        self.Ecinetique /= self.N
        self.energie = self.Ecinetique+self.Epotentielle/self.N

             
    def euler(self,h,hd2):
        for k in range(self.N):
            disque = self.liste_disques[k]
            self.deplacer_disque(disque,k,disque.x+h*disque.vx,disque.y+h*disque.vy)
            disque.vx += h*disque.ax
            disque.vy += h*disque.ay
        self.calculer_forces()
    def eulerA(self,h,hd2):
        for k in range(self.N):
            disque = self.liste_disques[k]
            self.deplacer_disque(disque,k,disque.x+h*disque.vx,disque.y+h*disque.vy)
        self.calculer_forces()
        for k in range(self.N):
            disque = self.liste_disques[k]
            disque.vx += h*disque.ax
            disque.vy += h*disque.ay
    def verlet(self,h,hd2):
        for k in range(self.N):
            disque = self.liste_disques[k]
            disque.vx += hd2*disque.ax
            disque.vy += hd2*disque.ay
            self.deplacer_disque(disque,k,disque.x+h*disque.vx,disque.y+h*disque.vy)
        self.calculer_forces()
        for k in range(self.N):
            disque = self.liste_disques[k]
            disque.vx += hd2*disque.ax
            disque.vy += hd2*disque.ay

             
    def integration(self,methode,h,n):
        hd2 = h/2
        for i in range(n):
            methode(h,hd2)
             
    def dessiner_disques(self,screen,echelle,couleur):
        for k in range(self.N):
            disque = self.liste_disques[k]
            pygame.draw.ellipse(screen,couleur,[(disque.x-self.rayon)*echelle,(disque.y-self.rayon)*echelle,self.rayon*2*echelle,self.rayon*2*echelle],2)