import numpy as np
from mpl_toolkits import mplot3d 
import matplotlib.pyplot as plt
import matplotlib.tri as mtri
from mpl_toolkits import mplot3d
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d import proj3d
import math
import pylab
from red_refine import red_refine
import scipy.sparse
import scipy.sparse.linalg
from scipy.sparse import csr_matrix
import timeit
import sys
#-----------------------------------------------------------------------
def get_geom():
    coord = np.asarray([[0,0],[0.5,0],[1,0],[1,0.5],[1,1],[0.5,1],
                    [0,1],[0,0.5],[0.25,0.25],[0.75, 0.25],[0.75,0.75],
                    [0.25,0.75],[0.5,0.5]])
    triangles = np.asarray([[0,1,8],[1,12,8],[1,9,12],[1,2,9],[2,3,9],
                     [3,12,9],[3, 10, 12],[3,4,10], [4,5,10],
                     [5,12,10],[5,11,12],[5,6,11],[6,7,11],[7,12,11],
                     [7,8,12],[7,0,8]])
    dirichlet= np.array([[0,1],[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,0]])
    neumann= np.zeros([0, 2])
    return coord, triangles, dirichlet, neumann

def restrict2dof(dof,nnodes):
    ndof=np.size(dof)
    R=csr_matrix((np.ones(ndof),(dof,np.arange(0,ndof))),shape = (nnodes,ndof))
    return R

def stiffness_matrix(coord,triangles):
    nelems=np.size(triangles,0)
    nnodes=np.size(coord,0)
    Alocal=np.zeros((nelems,3,3))
    I1=np.zeros((nelems,3,3))
    I2=np.zeros((nelems,3,3))

    for j in range(0,nelems):
        nodes_loc=triangles[j,:]
        coord_loc=coord[nodes_loc,:]
        T=np.array([coord_loc[1,:]-coord_loc[0,:] ,
               coord_loc[2,:]-coord_loc[0,:] ])
        area = 0.5 * ( T[0,0]*T[1,1] - T[0,1]*T[1,0] )
        tmp1= np.concatenate((np.array([[1,1,1]]), coord_loc.T),axis=0)
        tmp2= np.array([[0,0],[1,0],[0,1]])
        grads = np.linalg.solve(tmp1,tmp2)
        Alocal[j,:,:]=area* np.matmul(grads,grads.T)
        I1[j,:,:] = np.concatenate((np.array([nodes_loc]),np.array([nodes_loc]),np.array([nodes_loc])),axis=0)
        I2[j,:,:] = np.concatenate((np.array([nodes_loc]).T,np.array([nodes_loc]).T,np.array([nodes_loc]).T),axis=1)
        
    Alocal=np.reshape(Alocal,(9*nelems,1)).T
    I1=np.reshape(I1,(9*nelems,1)).T
    I2=np.reshape(I2,(9*nelems,1)).T
    A=csr_matrix((Alocal[0,:],(I1[0,:],I2[0,:])),shape = (nnodes,nnodes))
    return A

def RHS_vector(coord,triangles,f):
    nelems=np.size(triangles,0)
    nnodes=np.size(coord,0)
    b=np.zeros(nnodes)
    
    for j in range(0,nelems):
        nodes_loc=triangles[j,:]
        coord_loc=coord[nodes_loc,:]
        tmp=np.array([coord_loc[1,:]-coord_loc[0,:] ,
               coord_loc[2,:]-coord_loc[0,:] ])
        area = 0.5 * ( tmp[0,0]*tmp[1,1] - tmp[0,1]*tmp[1,0] )
        mid=1/3*(coord_loc[0,:]+coord_loc[1,:]+coord_loc[2,:])
        b[nodes_loc]=b[nodes_loc]+area/3*f(mid[0],mid[1])
        
    return b

def relax(A, b, u,n_smooth):
    for _ in range(n_smooth):
        u = u - 1/8 * (A*u-b)
    return u

def Wcycle(coarse,fine,A,b,x,n_smooth,Slist,Plist):
    nlevels=fine-coarse
    S=Slist[nlevels]
    P=Plist[nlevels]
    A_inner=(S.transpose()@A)@S
    b_inner=S.transpose()@b
    if coarse==fine:
        x=S*scipy.sparse.linalg.spsolve(A_inner,b_inner)
    else:
        for _ in range(0,2):
            x=S*relax(A_inner, b_inner, S.transpose()*x,n_smooth)
            r=b-A*x;
            q=Wcycle(coarse,fine-1,P.transpose()@(A@P),P.transpose()*r,0*P.transpose()*r,n_smooth,Slist,Plist)
            x=x+P*q
            x=S*relax(A_inner, b_inner, S.transpose()*x,n_smooth)
    return x
    
def FEM_mg(coarse,fine,A,b,n_iter,n_smooth,Slist,Plist,x):  
    for m in range(0,n_iter):
        x = Wcycle(coarse,fine,A,b,x,n_smooth,Slist,Plist)
    return x

#-------- the numerical experiment --------------------------3
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

u_exact=np.vectorize(uex_fun)
f=np.vectorize(f_fun)

coord, triangles, dirichlet, neumann = get_geom()
Slist=[]
Plist=[]
coarse=1
fine=6
n_iter=2
n_smooth=2
n_meshes=fine-coarse+1
ndoflist=np.zeros(n_meshes)
errlist=np.zeros(n_meshes)
for j in range(0,n_meshes):
    lvl=coarse+j
    print('mesh hierarchy and assembly...')
    coord, triangles, dirichlet,_,_,P = \
           red_refine(coord, triangles, dirichlet, neumann)
    Plist.append(P)
    nnodes=np.size(coord,0)
    dbnodes=np.unique(dirichlet)
    dof=np.setdiff1d(range(0,nnodes),dbnodes)
    ndof=np.size(dof)
    Slist.append(restrict2dof(dof,nnodes))
    A=stiffness_matrix(coord,triangles)
    b=RHS_vector(coord,triangles,f)
    print('done')
    start = timeit.default_timer()
    if j==0:x_init=0*b
    else:x_init=P@x
    x=FEM_mg(coarse,lvl,A,b,n_iter,n_smooth,Slist,Plist,x_init)
    S=Slist[lvl-1]
    x_ref=u_exact(coord[:,0],coord[:,1])
    stop = timeit.default_timer()
    e=x-x_ref.transpose()
    enerr=np.sqrt(e.transpose()@A@e)
    ndoflist[j]=ndof
    errlist[j]=enerr
    print('mesh ',j+1,'\b/',n_meshes,'',
          '  ndof=',"%10.3E" %ndof,'  energy error=',"%10.3E" %enerr,'  solve time: ', "%10.3E" %(stop - start))

#-----------------------------------------------------------------------
# Visualization
#-----------------------------------------------------------------------
# convergence history
fig = plt.figure()
ax = plt.axes()
ax.loglog(ndoflist,errlist,label='error',marker='x')
ax.loglog(ndoflist,errlist[0]/ndoflist[0]*ndoflist**(-1/2),label='slope -1/2',linestyle='dotted')
ax.set_xlabel('ndof')
ax.set_title('energy error I_h(u)-u_{multigrid}')
ax.legend()
plt.tight_layout()
plt.show()