import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
from scipy.spatial import Delaunay
from scipy.sparse import lil_matrix
from scipy.sparse.linalg import spsolve

def generate_mesh(n, orientation="null"): 
    x = np.linspace(0, 1, n)
    y = np.linspace(0, 1, n)
    X, Y = np.meshgrid(x, y)
    points = np.c_[X.ravel(), Y.ravel()]
    simplices = []
    for i in range(n - 1):
        for j in range(n - 1):
            # Indices des coins d'un rectangle
            p0 = i * n + j  # Bas gauche
            p1 = i * n + (j + 1)  # Bas droite
            p2 = (i + 1) * n + j  # Haut gauche
            p3 = (i + 1) * n + (j + 1)  # Haut droite
            
            if orientation == "left":
                # Division selon la diagonale gauche
                simplices.append([p0, p2, p3])
                simplices.append([p0, p3, p1])
            elif orientation == "right":
                # Division selon la diagonale droite
                simplices.append([p0, p2, p1])
                simplices.append([p2, p3, p1])
            else:  # Orientation "null" : Delaunay (par défaut)
                simplices.extend(Delaunay(points).simplices)
                return points, simplices
    
    return points, np.array(simplices)


def assemble_stiffness_matrix(points, simplices):
    """
    Assemble la matrice de raideur pour -Δu + u avec FEM.
    """
    num_points = len(points)
    A = lil_matrix((num_points, num_points)) 
    
    for tri in simplices:
        p0, p1, p2 = points[tri]
        J = np.array([[p1[0] - p0[0], p2[0] - p0[0]], 
                    [p1[1] - p0[1], p2[1] - p0[1]]])
        detJ = np.abs(np.linalg.det(J))
        B = np.linalg.inv(J).T @ np.array([[-1, 1, 0], [-1, 0, 1]])
        stiffness_local = (B.T @ B) * detJ / 2 + np.eye(3) * detJ / 12
        for i in range(3):
            for j in range(3):
                A[tri[i], tri[j]] += stiffness_local[i, j]
    
    return A

def assemble_load_vector(points, simplices, f_func):
    """
    Assemble le vecteur de charges pour f(x, y).
    """
    num_points = len(points)
    b = np.zeros(num_points)
    
    for tri in simplices:
        p0, p1, p2 = points[tri]
        detJ = np.abs(np.linalg.det(np.array([[p1[0] - p0[0], p2[0] - p0[0]], 
                                            [p1[1] - p0[1], p2[1] - p0[1]]])))
        f_avg = f_func((p0[0] + p1[0] + p2[0]) / 3, 
                    (p0[1] + p1[1] + p2[1]) / 3)
        b_local = f_avg * detJ / 6
        for i in range(3):
            b[tri[i]] += b_local
    
    return b

def apply_dirichlet_bc(A, b, points, boundary_value=0):
    """
    Applique les conditions de Dirichlet homogènes u = 0 sur le bord.
    """
    for i, (x, y) in enumerate(points):
        if x == 0 or x == 1 or y == 0 or y == 1:
            A[i, :] = 0
            A[i, i] = 1
            b[i] = boundary_value

def solve_fem(n, f_expr):
    """
    Résout -Δu + u = f sur le domaine d'étude.
    """
    x, y = sp.symbols('x y')
    f_func = sp.lambdify((x, y), f_expr, 'numpy')
    points, simplices = generate_mesh(n, "right")
    A = assemble_stiffness_matrix(points, simplices)
    b = assemble_load_vector(points, simplices, f_func)
    apply_dirichlet_bc(A, b, points)
    u = spsolve(A.tocsr(), b)
    return points, simplices, u


# f(x, y) = sin(πx)sin(πy)
x, y = sp.symbols('x y')
f = sp.sin(sp.pi * x) * sp.sin(sp.pi * y)

# Résolution
n = 50
points, simplices, u = solve_fem(n, f)

# plt.triplot(points[:, 0], points[:, 1], simplices)
# plt.title("Maillage")
# plt.xlabel('x')
# plt.ylabel('y')
# plt.show()

fig = plt.figure(figsize=(10, 12))
ax = fig.add_subplot(111, projection='3d')
ax.plot_trisurf(points[:, 0], points[:, 1], simplices, u, cmap='viridis', edgecolor='k', linewidth=0.2)
ax.set_title("-Δu + u = f")
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('u(x, y)')
plt.show()