from pylab import *

### Question 1) 

def Voisins(M,i,j) :
    p,q=M.shape[:2]
    L=[]
    if i>0 :
        L.append([i-1,j])
    if i<p-1 :
        L.append([i+1,j])
    if j>0 :
        L.append([i,j-1])
    if j<q-1 :
        L.append([i,j+1])
    return L




def Dynamique(M_0,M_1,alpha) :
    """ calcule l'état suivant, à partir de deux états conécutifs """
    p,q=M_0.shape
    new_M=zeros((p,q))
    for i in range(p) :
        for j in range(q) :
            c=0
            s=0
            for v in Voisins(new_M,i,j) :
                c+=1
                s+=M_1[v[0],v[1]]
            s*=alpha
            s+=(2-alpha*c)*M_1[i,j]
            new_M[i,j]=s-M_0[i,j]
    return new_M

### Question 2) 

def Diffusion(M_0,M_1,alpha) :
    P=M_0.copy()
    Q=M_1.copy()
    c=0
    ion() # mode interactif
    figure(1) # nouvelle figure
    axis('off')
    image=imshow(Q,cmap="jet",interpolation="nearest") #trace de l'image
    colorbar()
    show() # affichage
    pause(0.001) # necessaire (bug matplotlib)
    while fignum_exists(1)  : # tant qu'on n'a pas detruit la figure
        M=Dynamique(P,Q,alpha).copy()
        P=Q.copy()
        Q=M.copy()
        image.set_data(M) # on met a jour les donnees de la figure
        title("Propagation de l'onde après "+str(c)+" étapes.")
        axis('off')
        image.changed() # on indique que l'image a change, il faudra la redessiner
        draw() # on redessine
        pause(0.001) # on attend un peu
        c+=1
        show() # necessaire (bug matplotlib)
    ioff() # fin du mode interactif
def Diffusion_Bis(M_0,M_1,alpha) :
    """ légèrement plus rapide que Diffusion car évite l'appel aux fonctions et la création de matrices intermédiaires """
    
    P=M_0.copy()
    Q=M_1.copy()
    p,q=P.shape
    ion() # mode interactif
    figure(1) # nouvelle figure
    image=imshow(Q,cmap="jet",interpolation="nearest") #trace de l'image
    colorbar()
    show() # affichage
    pause(0.001) # necessaire (bug matplotlib)
    new_M=zeros((p,q))
    compteur=0
    while fignum_exists(1)  : # tant qu'on n'a pas detruit la figure
        compteur+=1
        for i in range(p) :
            for j in range(q) :
                c=0
                s=0
                for v in Voisins(new_M,i,j) :
                    c+=1
                    s+=Q[v[0],v[1]]
                s*=alpha
                s+=(2-alpha*c)*Q[i,j]
                new_M[i,j]=s-P[i,j]
        P=Q.copy()
        Q=new_M.copy()
        
        image.set_data(new_M) # on met a jour les donnees de la figure
        title("Propagation de l'onde après "+str(compteur)+" étapes.")
        image.changed() # on indique que l'image a change, il faudra la redessiner
        draw() # on redessine
        pause(0.001) # on attend un peu
        
        c+=1
        ioff() # fin du mode interactif
        show() # necessaire (bug matplotlib)
### Question 3) 

def Conditions_Initiales(p,q,r_int,r_ext,profondeur,c_x,c_y) :
    """ renvoie la condition initiale [M_0,M_1] de taille p*q, avec une profondeur sur le disque de [c_x,c_y] et de rayon r_int et en compensation sur la couronne de rayon r_ext """
    M_0=zeros((p,q))
    M_1=zeros((p,q))
    X,Y=ogrid[0:p,0:q]
    mask_profond=(X-c_x)**2+(Y-c_y)**2<=r_int**2
    M_1[mask_profond]=-profondeur
    M_2=zeros((p,q))
    mask_couronne=((X-c_x)**2+(Y-c_y)**2>r_int**2) & ((X-c_x)**2+(Y-c_y)**2<=r_ext**2)
    M_2[mask_couronne]=1
    c=sum([sum(M_2[:,j]) for j in range(q)])
    M_3=zeros((p,q))
    M_3[mask_profond]=1
    d=sum([sum(M_3[:,j]) for j in range(q)])
    h=profondeur*d/c
    M_1[mask_couronne]=h
    return [M_0,M_1]
  
### Question 4) 
# 
# p,q=60,100
# r_int=2
# r_ext=4
# profondeur=10
# c_x=20
# c_y=30
# alpha=0.5
# 
# P,Q=Conditions_Initiales(p,q,r_int,r_ext,profondeur,c_x,c_y)
# Diffusion(P,Q,alpha)

def Verif(M) :
    """ calcule la somme de tous les coefficients """
    return sum([sum(M[:,j]) for j in range(M.shape[1])])
# print(Verif(P))
# for k in range(100) :
#     M=Dynamique(P,Q,alpha).copy()
#     P=Q.copy()
#     Q=M.copy()
#     print(Verif(Q))

### Question 5)

def New_Cond_Init(p,q,r_int,r_ext,profondeur,L_Pos) :
    """ L_Pos contient la liste des coordonnées des cailloux """
    M_0=zeros((p,q))
    M_1=zeros((p,q))
    
    def Nbre_1(i,j,r_int,L_Pos) :
        """ calcule le nombre de cailloux impactant [i,j] en profondeur """
        return len([c for c in L_Pos if ( (c[0]-i)**2+(c[1]-j)**2 <= r_int**2 )])
    
    def Nbre_2(i,j,r_int,r_ext,L_Pos) :
        """ calcule le nombre de cailloux impactant [i,j] en crête """
        return len([c for c in L_Pos if ( (c[0]-i)**2+(c[1]-j)**2 > r_int**2 and (c[0]-i)**2+(c[1]-j)**2 <= r_ext**2)])
    
    M_2=zeros((p,q,2))  
    
    for i in range(p) :
        for j in range(q) :
            M_2[i,j]=[Nbre_1(i,j,r_int,L_Pos),Nbre_2(i,j,r_int,r_ext,L_Pos)]
    
    Total_profond=sum([sum(M_2[:,j,0]) for j in range(q)])
    Total_crete=sum([sum(M_2[:,j,1]) for j in range(q)])
    coeff=profondeur*Total_profond/Total_crete
    
    for i in range(p) :
        for j in range(q) :
            M_1[i,j]=-M_2[i,j,0]*profondeur+coeff*M_2[i,j,1]
    return [M_0,M_1]
          

p,q=60,100
r_int=2
r_ext=4
profondeur=10
alpha=0.5

from random import *
N=1
L_Pos=[[randint(0,p-1),randint(0,q-1)] for k in range(N)]

P,Q=New_Cond_Init(p,q,r_int,r_ext,profondeur,L_Pos)
# Diffusion(P,Q,alpha)
Diffusion_Bis(P,Q,alpha)
# 
# print(Verif(P))
# for k in range(100) :
#     M=Dynamique(P,Q,alpha).copy()
#     P=Q.copy()
#     Q=M.copy()
#     print(Verif(Q))