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 FDM_coords(n):
    N=(n+1)**2
    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')
    return coord_x, coord_y

def FDM_data(n,f):
    h=1/n
    N=(n+1)**2
    data = np.array([-np.ones(N),-np.ones(N),4*np.ones(N),-np.ones(N),-np.ones(N)])
    diags = np.array([-n-1,-1,0, 1,n+1])
    A=spdiags(data, diags, N, N).tocsr();
    coord_x, coord_y=FDM_coords(n)
    b=h**2*np.reshape(f(coord_x,coord_y),(n+1)**2)
    return A, b, h

def restrict2dof(n):
    N=(n+1)**2
    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))
    dof=np.setdiff1d(range(0,N),bdry)
    ndof=np.size(dof)
    R=csr_matrix((np.ones(ndof),(dof,np.arange(0,ndof))),shape = (N,ndof))
    return R
    
def solveFDM(n,f):
    A, b, h = FDM_data(n,f)
    R=restrict2dof(n)
    A_inner=(R.transpose()@A)@R
    b_inner=R.transpose()@b
    x=R@scipy.sparse.linalg.spsolve(A_inner,b_inner)
    return x, h

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

coarse=2 #2^coarse intervals per axis
fine=7
n_meshes=fine-coarse+1
hlist=np.zeros(n_meshes)
errlist=np.zeros(n_meshes)
for j in range(0,n_meshes):
    lvl=coarse+j
    n=2**(lvl)
    x, h=solveFDM(n,f)
    # compute errors
    coord_x, coord_y=FDM_coords(n)
    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('mesh ',j+1,'\b/',n_meshes,'',"%10.3E" %(n+1)**2,' nodes',
          '  h=',"%10.3E" %h,'  max error=',"%10.3E" %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,label='slope -1',linestyle='dashed')
ax.loglog(1./hlist,errlist[0]/hlist[0]*hlist**(2/3),label='slope -2/3',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('')

plt.show()