import numpy as np
from numpy import matlib
import matplotlib.pyplot as plt
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.exp(x)*x*y*(1-x)*(1-y)
    return val

# right-hand side function f
def f_fun(x,y):
    val = -np.exp(x)*x*(x*(y**2-y+2)+3*y**2-3*y-2)
    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

# application of the SSOR preconditionoer
def applyC(omega,x):
    m=np.int_(np.sqrt(x.size))#number of inner nodes per axis
    #forward substitution
    a=np.zeros(m**2)
    a[0]=x[0]/4
    for j in range(1,m):
        a[j]=(x[j]+omega*a[j-1])/4
    for j in range(1,m):
        ind=j*m
        a[ind]=(x[ind]+omega*a[ind-m])/4
        for k in range(1,m):
            ind=j*m+k
            a[ind]=(x[ind]+omega*a[ind-1]+omega*a[ind-m])/4            
    #backward substitution
    x=np.zeros(m**2)
    x[m**2-1]=a[m**2-1]/4
    for j in range(1,m):
        ind=m**2-j-1
        x[ind]=(a[ind]+omega*x[ind+1])/4
    for j in range(1,m):
        ind=m**2-j*m-1
        x[ind]=(a[ind]+omega*x[ind+m])/4
        for k in range(1,m):
            ind=m**2-j*m-k-1
            x[ind]=(a[ind]+omega*x[ind+1]+omega*x[ind+m])/4
    return x

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

n_iter=100
cg_residuals=np.zeros(n_iter)
pcg_residuals=np.zeros(n_iter)

lvl=6
n=2**(lvl)    
A, b, h = FDM_data(n,f)
S=restrict2dof(n)
A=(S.transpose()@A)@S
b=S.transpose()@b

#pcg
omega=1.3
x=np.zeros(b.size)
normb=np.sqrt(np.dot(b,b))
g=A*x-b
d=-applyC(omega,g)
tmp=-np.dot(g,d)
for m in range(0,n_iter):
    alpha=tmp/(np.dot(d,A*d))
    x=x+alpha*d
    g=g+alpha*A*d
    Cg=applyC(omega,g)
    tmp2=np.dot(g,Cg)
    beta=tmp2/tmp
    d=-Cg+beta*d
    tmp=tmp2
    pcg_residuals[m]=np.sqrt(np.dot(g,g))/normb

#cg
x=np.zeros(b.size)
g=A*x-b
d=-g
tmp=-np.dot(g,d)
for m in range(0,n_iter):
    alpha=tmp/(np.dot(d,A*d))
    x=x+alpha*d
    g=g+alpha*A*d
    tmp2=np.dot(g,g)
    beta=tmp2/tmp
    d=-g+beta*d
    tmp=tmp2
    cg_residuals[m]=np.sqrt(np.dot(g,g))/normb

#-----------------------------------------------------------------------
# Visualization
#-----------------------------------------------------------------------
# convergence history
fig = plt.figure()
ax = plt.axes()
ax.semilogy(range(0,n_iter),pcg_residuals,label='pcg residuals SSOR omega='+str(omega),marker='x')
ax.semilogy(range(0,n_iter),cg_residuals,label='cg residuals',marker='x')
ax.set_xlabel('iteration')
ax.set_ylabel('relative residual')
R=str(n+1)
ax.set_title('comparison pcg vs cg for FDM system with '+R+'x'+R+' nodes')
ax.legend()
plt.tight_layout()
plt.show()