import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
import matplotlib.tri as mtri

# solves Poisson equation Lu = f for a given triangulation, L = -Laplacian
# here supplied with simple uniform triangulation of unit square
# homogeneous Dirichlet boundary conditions applied

def g(x):
    return 10*np.sin(np.pi*x[0])*np.cos(np.pi*x[1])

def uniform_grid(n):
    # generate a set of points
    x,y=np.mgrid[0:n+1.,0:n+1.]/n
    p = np.vstack([x.flatten(),y.flatten()])
    # construct triangulation of the points
    # this is not a good method for generating the uniform grid here!
    # it is used because p can be readily replaced by any set of points
    tri = sp.spatial.Delaunay(p.T)
    return tri

def assemble(tri):
    N = tri.simplices.shape[0]
    M,d = tri.points.shape
    # initialize
    A = np.zeros((M,M))
    f = np.zeros((M,1))
    # used to construct stiffness
    Phi = np.hstack([np.eye(d),-np.ones((d,1))])
    lbda = np.ones((d+1,1))/(d+1)
    for i in range(N):
        # collect vertex numbers
        ind = tri.simplices[i]
        # lookup vertex coordinates
        X = tri.points[ind,:]
        # Jacobian of forward mapping (uses transpose)
        J = X[0:d]-X[d]
        Jac=np.absolute(np.linalg.det(J))
        # Gradient matrix 
        # Note that J and Phi are transposes of what appears in the notes and
        # the matlab code, this is because python does not have a
        # 'matrix right division' like matlab
        G = np.linalg.solve(J,Phi).T
        # element stiffness
        A1 = (G.dot(G.T)).dot(Jac)/sp.math.factorial(d)
        # element load by 1-point barycentric quadrature
        bary = (X.T).dot(lbda)
        f1 = g(bary)*lbda*Jac/sp.math.factorial(d)
        # insert element stiffness/load into main
        f_ind=np.ix_(ind,ind)
        A[f_ind] = A[f_ind] + A1
        f[ind] = f[ind] + f1
    return A,f

def main(n):
    # generate triangulation
    tri = uniform_grid(n)
    # assemble stiffness and load
    A,f=assemble(tri)
    # implement boundary conditions
    boundary=np.unique(tri.convex_hull)
    A[boundary]=0
    f[boundary]=0
    for j in boundary:
        A[j,j]=1
    # solve system
    u=np.linalg.solve(A,f)
    # plot
    triang = mtri.Triangulation(tri.points[:][:,0],tri.points[:][:,1],triangles=tri.simplices)
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.plot_trisurf(triang, u.ravel(), cmap=plt.cm.CMRmap)
    plt.show()
    
main(30)