import matplotlib.pyplot as plt
import math
import numpy as np
import scipy as scp
import numpy.random as rd
import time

#Généralités pour décrire le domaine choisi
def graphe_bords(B):
    """Affiche un graphique des bords du domaine.
    """
    Nx,Ny = B.shape
    plt.figure()
    plt.imshow(B,origin='lower',cmap='binary', interpolation='Nearest')
    plt.colorbar()
    plt.title('matrice B des bords du domaine')
    plt.show()

def graphe_equipot(B,f,a=0,nb_equipot=25):
    """Affiche les bords du domaine B en noir, et des surfaces équipotentielles de f."""
    Nx,Ny = B.shape
    plt.figure()
    x = np.linspace(0,Nx-1,Nx)        # de 0 à Nx-1 inclus
    y = np.linspace(0,Ny-1,Ny)
    plt.imshow(B,origin='lower',cmap='binary',interpolation='nearest')
    if (a==1):
        plt.imshow(f,origin='lower')  # colorplot de f
        plt.colorbar()
    cont=plt.contour(y,x,f,nb_equipot,colors='k')
    plt.title('equipotentielles')
    plt.clabel(cont,fmt='%1.1f') # le 1.1 controle le nombre de chiffres après la virgule sur l'affichage
    plt.show()

def w_opt(Nx,Ny):
    return 2./(1.+np.pi/(Nx*Ny)*((Nx**2+Ny**2)/2.)**0.5)
#
def Magnus_():

    def initialisation_effet_magnus(B,phi,a,b,R):
        """Initialise un cadre avec à gauche phi=a, à droite phi=b, en bas et en haut
        phi qui va linéairement de a jusqu'a b, et un cercle représentant le cylindre, au potentiel a."""
        # bord bas :
        B[0,:] = 1
        phi[0:] = np.linspace(a,b,Ny)
        # bord haut :
        B[0,:] = 1
        phi[0:] = np.linspace(a,b,Ny)
        # bord gauche :
        B[:,Ny-1] = 1
        phi[:,Ny-1] = a
        # bord droit :
        B[:,Ny-1] = 1
        phi[:,Ny-1] = b
        # Cercle centré au centre de l'image :
        i0 = Nx//2
        j0 = Ny//2
        for i in range(Nx):
            for j in range(Ny):
                if (i-i0)**2+(j-j0)**2 <= (R)**2:  # cercle de centre i0, j0 et de rayon R
                    B[i,j] = 1
                    phi[i,j] = a

    def une_iteration_GS(B,f):
        f2 = f.copy()
        ecart = 0.
        for i in range(1,Nx-1):  # on va de 1 à N-2 car pas traiter le cadre (Dirichlet)
            for j in range(1,Ny-1):
                if B[i,j]==0:    #pas sur un bord
                    f[i,j] = (1.-w)*f[i,j] + w*(f[i+1,j]+f[i,j+1]+f[i-1,j]+f[i,j-1])/4.
                    ecart = ecart + (f[i,j] - f2[i,j])**2
        return (ecart/(Nx*Ny))**0.5

    def iterations_GS(B,f,eps):
        ecart = eps+1.
        nb_iterations = 0
        while ecart > eps:
            ecart = une_iteration_GS(B,f)
            nb_iterations += 1
        return nb_iterations

    def calcul_V(phi):
        Vx = np.zeros((Nx,Ny))
        Vy = np.zeros((Nx,Ny))
        norme_V = np.zeros((Nx,Ny))
        for i in range(1,Nx-1):
            for j in range(1,Ny-1):
                Vx[i,j] = -(phi[i+1,j]-phi[i-1,j])/(2.*delta)
                Vy[i,j] = -(phi[i,j+1]-phi[i,j-1])/(2.*delta)
        norme_V = (Vx**2+Vy**2)**0.5

        return Vx, Vy, norme_V

    Nx, Ny = 150,150
    w = w_opt(Nx,Ny)
    delta = 3.e-2 # pas de la grille en m
    B = np.zeros((Nx,Ny),bool)
    phi = np.zeros((Nx,Ny))
    initialisation_effet_magnus(B,phi,0,(Nx-1)*delta*1.2,0.055)
    iterations_GS(B,phi,1.e-3)

    Vx, Vy, norme_V = calcul_V(phi)
    max_V = np.max(norme_V)
    print("Valeur maximale de V : " + str(max_V) + "m/s")
    graphe_equipot(B,phi,1,40)