import numpy as np
import scipy.sparse as sparse
import scipy.sparse.linalg as linalg
from scipy.linalg import solve

DIR = 1
NEU = 2

class Noeud:
    def __init__(self,x,y,a,s,lim=0,ksi=0,K0=0,u0=0,g0=0,bord=False):
        self.xy = np.array([x,y])
        self.x = x 
        self.y = y
        self.a = a
        self.s = s
        self.lim = lim
        self.ksi = ksi
        self.K0 = K0
        self.u0 = u0
        self.g0 = g0
        self.noeuds_voisins = []
        self.indice = 0 # indice dans le maillage
        self.bord= bord
    def ajouterVoisin(self,noeud):
        self.noeuds_voisins.append(noeud)              
    def extraireTriangles(self):
        N = len(self.noeuds_voisins)
        triangles =  []
        for i in range(N-1):
            n1 = self
            n2 = self.noeuds_voisins[i]
            n3 = self.noeuds_voisins[i+1]
            triangles.append([n1,n2,n3])
        if not(self.bord):
            n1 = self
            n2 = self.noeuds_voisins[N-1]
            n3 = self.noeuds_voisins[0]
            triangles.append([n1,n2,n3])   
            
        return triangles
                    

class ElementsFinis2D:
    def __init__(self):
        self.noeuds_dirichlet = []
        self.noeuds = []
        self.indice = 0
        self.indice_dirichlet = 0
    def ajouterNoeud(self,noeud):
        if noeud.lim == DIR:
            noeud.indice = self.indice_dirichlet
            self.noeuds_dirichlet.append(noeud)
            self.indice_dirichlet += 1
        else:
            noeud.indice = self.indice
            self.noeuds.append(noeud)
            self.indice += 1
    def definirMaillage(self,maillage):
        self.noeuds_dirichlet = []
        self.noeuds = []
        self.indice = 0
        self.indice_dirichlet = 0
        for noeud in maillage:
            self.ajouterNoeud(noeud)
    def fonctionPhi(self,triangle):
        n1 = triangle[0]
        n2 = triangle[1]
        n3 = triangle[2]
        detT = n2.x*n3.y-n2.x*n1.y-n1.x*n3.y-n2.y*n3.x+n2.y*n1.x+n1.y*n3.x
        if detT <=0 : 
            print("Erreur de définition du triangle")
            print(n1.xy)
            print(n2.xy)
            print(n3.xy)
        ST = 0.5*detT
        c0 = (n2.y-n3.y)/detT
        c1 = (n3.x-n2.x)/detT 
        c2 = (n2.x*n3.y-n3.x*n2.y)/detT
        return [ST,c0,c1,c2]
    def creerMatrices(self):
        self.Nnoeuds = len(self.noeuds)
        self.M = np.zeros((self.Nnoeuds,self.Nnoeuds),dtype=np.float64)
        self.B = np.zeros(self.Nnoeuds,dtype=np.float64)
    def matriceA(self):
        for i in range(self.Nnoeuds):
            Ni = self.noeuds[i]
            triangles = Ni.extraireTriangles()
            Aii = 0
            for t in range(len(triangles)):
                [ST,ci0,ci1,ci2] = self.fonctionPhi(triangles[t])
                Aii += (triangles[t][0].a+triangles[t][1].a+triangles[t][2].a)/3*(ci0*ci0+ci1*ci1)*ST
                if triangles[t][1].lim != DIR:
                    j = triangles[t][1].indice
                    [ST,cj0,cj1,cj2] = self.fonctionPhi([triangles[t][1],triangles[t][2],triangles[t][0]])
                    self.M[i,j] += (triangles[t][0].a+triangles[t][1].a+triangles[t][2].a)/3*(ci0*cj0+ci1*cj1)*ST
                if triangles[t][2].lim != DIR:
                    j = triangles[t][2].indice
                    [ST,cj0,cj1,cj2] = self.fonctionPhi([triangles[t][2],triangles[t][0],triangles[t][1]])
                    self.M[i,j] += (triangles[t][0].a+triangles[t][1].a+triangles[t][2].a)/3*(ci0*cj0+ci1*cj1)*ST
            self.M[i,i] += Aii
    
    def matriceR(self):
        for i in range(self.Nnoeuds):
            Ni = self.noeuds[i]
            Rii = 0
            if Ni.lim==NEU:
                for n in Ni.noeuds_voisins:
                    if n.lim == NEU:
                        d = np.sqrt((Ni.x-n.x)**2+(Ni.y-n.y)**2)
                        Rii += -(Ni.a*Ni.K0+n.a*n.K0)/2*d/3
                        j = n.indice
                        self.M[i,j] += -1/6*(Ni.a*Ni.K0+n.a*n.K0)/2*d
            self.M[i,i] += Rii
    def matriceB1(self):
        for i in range(self.Nnoeuds):
            B1i = 0
            Ni = self.noeuds[i]
            triangles = Ni.extraireTriangles()
            for t in range(len(triangles)):
                [ST,ci0,ci1,ci2] = self.fonctionPhi(triangles[t])
                if triangles[t][1].lim == DIR :
                    [ST,cj0,cj1,cj2] = self.fonctionPhi([triangles[t][1],triangles[t][2],triangles[t][0]])
                    B1i += -triangles[t][1].ksi*(triangles[t][0].a+triangles[t][1].a+triangles[t][2].a)/3*(ci0*cj0+ci1*cj1)*ST
                if triangles[t][2].lim == DIR:
                    [ST,cj0,cj1,cj2] = self.fonctionPhi([triangles[t][2],triangles[t][0],triangles[t][1]])
                    B1i += -triangles[t][2].ksi*(triangles[t][0].a+triangles[t][1].a+triangles[t][2].a)/3*(ci0*cj0+ci1*cj1)*ST
            self.B[i] += B1i
    def matriceB2(self):
        for i in range(self.Nnoeuds):
            B2i = 0
            Ni = self.noeuds[i]
            triangles = Ni.extraireTriangles()
            for t in range(len(triangles)):
                [ST,ci0,ci1,ci2] = self.fonctionPhi(triangles[t])
                B2i += (triangles[t][0].s+triangles[t][1].s+triangles[t][2].s)/9*ST
            self.B[i] += B2i
    def matriceB3(self):
        for i in range(self.Nnoeuds):
            Ni = self.noeuds[i]
            if Ni.lim == NEU:
                B3i = 0
                for n in Ni.noeuds_voisins:
                    if n.lim == DIR:
                        d = np.sqrt((Ni.x-n.x)**2+(Ni.y-n.y)**2)
                        B3i += n.ksi*(Ni.a*Ni.K0+n.a*n.K0)/2*d/6
                self.B[i] += B3i
    def matriceB4(self):
        for i in range(self.Nnoeuds):
            Ni = self.noeuds[i]
            if Ni.lim == NEU:
                B4i = 0
                for n in Ni.noeuds_voisins:
                    if n.lim == NEU:
                        d = np.sqrt((Ni.x-n.x)**2+(Ni.y-n.y)**2)
                        B4i += 0.5*(Ni.a*(Ni.g0-Ni.K0*Ni.u0)+n.a*(n.g0-n.K0*n.u0))*0.5*d
                self.B[i] += B4i
    def solve(self,solver="iterative"):
        self.creerMatrices()
        self.matriceA()
        self.matriceR()
        self.matriceB1()
        self.matriceB2()
        self.matriceB3()
        self.matriceB4()
        if solver=="iterative":
            M = sparse.bsr_array(self.M)
            ksi,exitCode = linalg.cg(M, self.B, atol=1e-5)
            print("exitCode : ",exitCode)
        else:
            ksi = solve(self.M,self.B)
        for i in range(self.Nnoeuds):
            self.noeuds[i].ksi = ksi[i]
                    

N = 50
x = np.linspace(0,1,N)
y = np.linspace(0,1,N)
maillage = []
a = 1
s = 0
for j in range(N):
    for i in range(N):
        if i!=0 and i!=N-1 and j!=0 and j!=N-1:
            n = Noeud(x[i],y[j],a,s)
        else:  
            if i==0:
                n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
            elif i==N-1:
                n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=1,bord=True)
            if j==0 and i!=0 and i!=N-1:
                n = Noeud(x[i],y[j],a,s,lim=NEU,g0=0,bord=True)
            elif j==N-1 and i!=0 and i!=N-1:
                n = Noeud(x[i],y[j],a,s,lim=NEU,g0=0,bord=True)
        
        maillage.append(n)
        

def attribuerVoisins(N,maillage):
    indice = 0
    for j in range(N):
        for i in range(N):
            n = maillage[indice]
            if i!=0 and i!=N-1 and j!=0 and j!=N-1:
                n.ajouterVoisin(maillage[indice+1])
                n.ajouterVoisin(maillage[indice+N+1])
                n.ajouterVoisin(maillage[indice+N])
                n.ajouterVoisin(maillage[indice-1])
                n.ajouterVoisin(maillage[indice-1-N])
                n.ajouterVoisin(maillage[indice-N])
            else:
                if i==0 and j==0:
                    n.ajouterVoisin(maillage[indice+1])
                    n.ajouterVoisin(maillage[indice+N+1])
                    n.ajouterVoisin(maillage[indice+N])
                elif i==0 and j==N-1:
                    n.ajouterVoisin(maillage[indice-N])
                    n.ajouterVoisin(maillage[indice+1])
                elif i==0:
                    n.ajouterVoisin(maillage[indice-N])
                    n.ajouterVoisin(maillage[indice+1])
                    n.ajouterVoisin(maillage[indice+N+1])
                    n.ajouterVoisin(maillage[indice+N])
                if i==N-1 and j==0:
                    n.ajouterVoisin(maillage[indice+N])
                    n.ajouterVoisin(maillage[indice-1])
                elif i==N-1 and j==N-1:
                    n.ajouterVoisin(maillage[indice-1])
                    n.ajouterVoisin(maillage[indice-1-N])
                    n.ajouterVoisin(maillage[indice-N])
                elif i==N-1:
                    n.ajouterVoisin(maillage[indice+N])
                    n.ajouterVoisin(maillage[indice-1])
                    n.ajouterVoisin(maillage[indice-1-N])
                    n.ajouterVoisin(maillage[indice-N])
                if j==0 and i!=0 and i!=N-1:
                    n.ajouterVoisin(maillage[indice+1])
                    n.ajouterVoisin(maillage[indice+N+1])
                    n.ajouterVoisin(maillage[indice+N])
                    n.ajouterVoisin(maillage[indice-1])
                if j==N-1 and i!=0 and i!=N-1:
                    n.ajouterVoisin(maillage[indice-1])
                    n.ajouterVoisin(maillage[indice-N-1])
                    n.ajouterVoisin(maillage[indice-N])
                    n.ajouterVoisin(maillage[indice+1])
            indice+= 1
            
attribuerVoisins(N,maillage)
                    

elem = ElementsFinis2D()
elem.definirMaillage(maillage)
elem.solve()
U = np.zeros((N,N),dtype=np.float64)
indice = 0
for j in range(N):
    for i in range(N):
        n = maillage[indice]
        U[j,i] = n.ksi
        indice+= 1
from matplotlib.pyplot import *
figure()
contour(U,10,extent=(0,1,0,1))

                     

figure()
plot(x,U[N//2,:])
grid()
                     

N = 50
x = np.linspace(0,1,N)
y = np.linspace(0,1,N)
maillage = []
a = 1
s = 1
for j in range(N):
    for i in range(N):
        if i!=0 and i!=N-1 and j!=0 and j!=N-1:
            n = Noeud(x[i],y[j],a,s)
        else:  
            if i==0:
                n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
            elif i==N-1:
                n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
            if j==0 and i!=0 and i!=N-1:
                n = Noeud(x[i],y[j],a,s,lim=NEU,g0=0,bord=True)
            elif j==N-1 and i!=0 and i!=N-1:
                n = Noeud(x[i],y[j],a,s,lim=NEU,g0=0,bord=True)
        
        maillage.append(n)
attribuerVoisins(N,maillage)
        

elem = ElementsFinis2D()
elem.definirMaillage(maillage)
elem.solve()
U = np.zeros((N,N),dtype=np.float64)
indice = 0
for j in range(N):
    for i in range(N):
        n = maillage[indice]
        U[j,i] = n.ksi
        indice+= 1
from matplotlib.pyplot import *
figure()
contour(U,10,extent=(0,1,0,1))

                     

figure()
plot(x,U[N//2,:])
grid()
                     

N = 50
x = np.linspace(0,1,N)
y = np.linspace(0,1,N)
maillage = []
a = 1
s = 1
for j in range(N):
    for i in range(N):
        if i!=0 and i!=N-1 and j!=0 and j!=N-1:
            n = Noeud(x[i],y[j],a,s)
        else:  
            if i==0:
                n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
            elif i==N-1:
                n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
            if j==0 and i!=0 and i!=N-1:
                n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
            elif j==N-1 and i!=0 and i!=N-1:
                n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
        
        maillage.append(n)
attribuerVoisins(N,maillage)
        

elem = ElementsFinis2D()
elem.definirMaillage(maillage)
elem.solve()
U = np.zeros((N,N),dtype=np.float64)
indice = 0
for j in range(N):
    for i in range(N):
        n = maillage[indice]
        U[j,i] = n.ksi
        indice+= 1
from matplotlib.pyplot import *
figure()
contour(U,10,extent=(0,1,0,1))

                     

figure()
plot(x,U[N//2,:])
grid()
                     

N = 50
x = np.linspace(0,1,N)
y = np.linspace(0,1,N)
maillage = []
a1 = 1
a2 = 2
s = 0
for j in range(N):
    for i in range(N):
        if i < N//2 : a=a1
        else : a=a2
        if i!=0 and i!=N-1 and j!=0 and j!=N-1:
            n = Noeud(x[i],y[j],a,s)
        else:  
            if i==0:
                n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
            elif i==N-1:
                n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=1,bord=True)
            if j==0 and i!=0 and i!=N-1:
                n = Noeud(x[i],y[j],a,s,lim=NEU,g0=0,bord=True)
            elif j==N-1 and i!=0 and i!=N-1:
                n = Noeud(x[i],y[j],a,s,lim=NEU,g0=0,bord=True)
        
        maillage.append(n)
attribuerVoisins(N,maillage)
        

elem = ElementsFinis2D()
elem.definirMaillage(maillage)
elem.solve()
U = np.zeros((N,N),dtype=np.float64)
indice = 0
for j in range(N):
    for i in range(N):
        n = maillage[indice]
        U[j,i] = n.ksi
        indice+= 1
from matplotlib.pyplot import *
figure()
contour(U,10,extent=(0,1,0,1))

                     

figure()
plot(x,U[N//2,:])
grid()
                     

N = 50
x = np.linspace(0,1,N)
y = np.linspace(0,1,N)
maillage = []
a = 1
s = 0
# bords de l'objet
e = 5
i1 = N//2-e 
i2 = N//2+e 
j1 = N//2-e 
j2 = N//2+e
for j in range(N):
    for i in range(N):
        if (i==i1 and j1<=j<=j2) or (i==i2 and j1<=j<=j2) or (j==j1 and i1<=i<=i2) or (j==j2 and i1<=i<=i2):
            n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=1)
        else:
            if i!=0 and i!=N-1 and j!=0 and j!=N-1:
                n = Noeud(x[i],y[j],a,s)
            else:  
                if i==0:
                    n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
                elif i==N-1:
                    n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
                if j==0 and i!=0 and i!=N-1:
                    n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
                elif j==N-1 and i!=0 and i!=N-1:
                    n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
        maillage.append(n)
attribuerVoisins(N,maillage)
        

elem = ElementsFinis2D()
elem.definirMaillage(maillage)
elem.solve()
U = np.zeros((N,N),dtype=np.float64)
indice = 0
for j in range(N):
    for i in range(N):
        n = maillage[indice]
        U[j,i] = n.ksi
        indice+= 1
from matplotlib.pyplot import *
figure()
contour(U,10,extent=(0,1,0,1))

                     

figure()
plot(x,U[N//2,:])
grid()
                     

N = 50
x = np.linspace(0,1,N)
y = np.linspace(0,1,N)
maillage = []
a = 1
s = 0
for j in range(N):
    for i in range(N):
        if i!=0 and i!=N-1 and j!=0 and j!=N-1:
            n = Noeud(x[i],y[j],a,s)
        else:  
            if i==0:
                n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=0,bord=True)
            elif i==N-1:
                n = Noeud(x[i],y[j],a,s,lim=DIR,ksi=1,bord=True)
            if j==0 and i!=0 and i!=N-1:
                n = Noeud(x[i],y[j],a,s,lim=NEU,g0=0,K0=-1,u0=0,bord=True)
            elif j==N-1 and i!=0 and i!=N-1:
                n = Noeud(x[i],y[j],a,s,lim=NEU,g0=0,K0=-1,u0=0,bord=True)
        
        maillage.append(n)
attribuerVoisins(N,maillage)
        

elem = ElementsFinis2D()
elem.definirMaillage(maillage)
elem.solve()
U = np.zeros((N,N),dtype=np.float64)
indice = 0
for j in range(N):
    for i in range(N):
        n = maillage[indice]
        U[j,i] = n.ksi
        indice+= 1
from matplotlib.pyplot import *
figure()
contour(U,10,extent=(0,1,0,1))

                     

figure()
plot(x,U[N//2,:])
grid()