import numpy as np
from numpy import matlib
from mpl_toolkits import mplot3d 
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d import proj3d
import scipy.sparse
import scipy.sparse.linalg
from scipy.sparse import csr_matrix
from scipy.sparse import spdiags

# exact solution for purposes of comparison
def uex_fun(x,y):
    val= np.sin(np.pi*x)*np.sin(np.pi*y)
    return val

# right-hand side function f
def f_fun(x,y):
    val= 2*np.pi**2*np.sin(np.pi*x)*np.sin(np.pi*y)
    return val
def Laplacef_fun(x,y):
    val= 4*np.pi**4*np.sin(np.pi*x)*np.sin(np.pi*y)
    return val

#FDM for the square
def FDM(n,f,Laplacef):
    h=1/n;
    N=(n+1)**2
    I=np.ones(N)
    data = np.array([-I,-4*I,-I,-4*I,20*I,-4*I,-I,-4*I,-I])
    diags = np.array([-n-2,-n-1,-n,-1,0, 1,n,n+1,n+2])
    A=spdiags(data, diags, N, N).tocsr();
    coord_x=np.reshape(matlib.repmat(1/n*np.arange(0,n+1),1,n+1),(N,1),order='F')
    coord_y=np.reshape(matlib.repmat(1/n*np.arange(0,n+1),n+1,1),(N,1),order='F')
    #identify boundary nodes
    bottom=np.arange(0,n+1)
    top=np.arange(n*(n+1),N)
    left=(n+1)*np.arange(1,n)
    right=left+n
    bdry=np.concatenate((bottom,right,top,left))
    #degrees of freedom
    dof=np.setdiff1d(range(0,N),bdry)
    ndof=np.size(dof)
    
    #the FDM system
    b=6*h**2*f(coord_x,coord_y)-.5*h**4*Laplacef(coord_x,coord_y)
    #Restriction matrix
    #we could simply do
    #A_inner=A[np.ix_(dof,dof)]
    #in principle, but this is worse in terms of performance
    R=csr_matrix((np.ones(ndof),(dof,np.arange(0,ndof))),shape = (N,ndof))
    A_inner=(R.transpose()@A)@R
    b_inner=R.transpose()@b
        
    x=np.zeros(N)
    x[dof]=scipy.sparse.linalg.spsolve(A_inner,b_inner)
    return x, h, coord_x, coord_y


#-------- the numerical experiment --------------------------
u_exact=np.vectorize(uex_fun)
f=np.vectorize(f_fun)
Laplacef=np.vectorize(Laplacef_fun)

# n must be an even number, n+1 grid points per axis
n=2
n_meshes=5
hlist=np.zeros(n_meshes)
errlist=np.zeros(n_meshes)
for j in range(0,n_meshes):
    n=2*n
    x, h, coord_x, coord_y=FDM(n,f,Laplacef)
    uex4nodes=np.reshape(u_exact(coord_x,coord_y),(n+1)**2)
    e=uex4nodes-x
    maxerr = np.max(np.abs(e))
    hlist[j]=h
    errlist[j]=maxerr
    print('h=',h,'  max error=',maxerr)


# Visualization
# convergence history
fig = plt.figure()
ax = plt.axes()
ax.loglog(1./hlist,errlist,label='error',marker='x')
ax.loglog(1./hlist,errlist[0]/hlist[0]*hlist**2,label='slope -2',linestyle='dotted')
ax.loglog(1./hlist,errlist[0]/hlist[0]*hlist**4,label='slope -4',linestyle='dashed')
ax.set_xlabel('1/h')
ax.set_title('convergence history')
ax.legend()
plt.tight_layout()
    

# Creating dataset
xdata = np.linspace(-1,1,n+1)
ydata = np.linspace(-1,1,n+1)
X,Y = np.meshgrid(xdata,ydata)

#plot computed solution
fig = plt.figure()
ax3d = fig.add_subplot(2,2,1,projection='3d')
surf=ax3d.plot_surface(X, Y, np.reshape(x,(n+1,n+1)), cmap="viridis")
ax3d.set_title('FDM solution')
ax3d.set_xlabel('X')
ax3d.set_ylabel('Y')
ax3d.set_zlabel('')

#plot exact solution
#fig = plt.figure()
ax3d = fig.add_subplot(2,2,2,projection='3d')
surf=ax3d.plot_surface(X, Y, np.reshape(uex4nodes,(n+1,n+1)), cmap="viridis")
ax3d.set_title('Exact solution')
ax3d.set_xlabel('X')
ax3d.set_ylabel('Y')
ax3d.set_zlabel('')

#plot computed solution
ax3d = fig.add_subplot(2,2,3,projection='3d')
surf=ax3d.plot_surface(X, Y, np.reshape(e,(n+1,n+1)), cmap="viridis")
ax3d.set_title('Error')
ax3d.set_xlabel('X')
ax3d.set_ylabel('Y')
ax3d.set_zlabel('')
plt.show()