#! /usr/bin/env python
#
import numpy as np
import scipy.linalg as la
'''
  Purpose:

    FEM1D_CLASSES defines classes for the finite element method.

  Discussion:

    The structure and interfaces of this code are meant to suggest
    the way in which the FEniCS program works.

    The classes defined here include:
    * Mesh
    * Shapefns
    * FiniteElement
    * FunctionSpace

  Modified:

    21 August 2014

  Author:

    Mike Sussman, Mathematics Department, University of Pittsburgh

{@(#) Sun May 18 17:31:02 2014 }
'''

class Mesh(object):
    '''
    A mesh is a list of point global coordinates
    and a list of element definitions (by endpoint)
    This class defines a uniform mesh on an interval
    Mesh(N,a,b) N intervals on [a,b]
    coordinates(n=None): return coord node n, or all coords
    cells(n=None): array of n-th cell's endpoint numbers, or all
    size(): returns number of elements
    '''

    def __init__(self,N,a,b):
        '''
        N is the number of elements = number of INTERVALS
        a and b are interval endpoints
        '''
        #  Define N elements using N+1 evenly spaced nodes.
        self.__nodes = np.linspace ( a, b, N + 1 )

        #  Cell I uses nodes I and I+1.
        self.__cells = np.array ( [ range(N), range(1,N+1) ] )

        self.__N = N

    def size ( self ):
      return self.__N

    def coordinates(self,n=None):

      if not ( n is None ):
        return self.__nodes[n]
      else:
        return self.__nodes

    def cells(self,n=None):

      if not ( n is None ):
        return self.__cells[:,n]
      else:
        return self.__cells


class Shapefns(object):
    '''
    Define Quadratic Lagrange shape functions
    These will be defined on the (local) interval [0,1], with
        mid point 0.5
    Shapefns()
    eval(n,xi): phi[n](xi)
    ddx(n,xi):  dphi[n](xi)
    size(): number of nodes for these shape functions
    '''

    def __init__(self):
        ''' 
        an array of functions for phi and deriv phi
        '''
        self.__phi = [lambda xi: 2.0 * (xi - 0.5) * (xi - 1.0) ,\
                     lambda xi: 4.0 * xi * (1.0 - xi), \
                     lambda xi: 2.0 * xi * (xi - 0.5)]
        # and dphi (derivative of phi w.r.t. xi)
        # derivative of second factor * first + derivative of first factor * second
        self.__dphi = [lambda xi: 2.0 * (xi - 0.5) + 2.0*(xi - 1.0) ,\
                      lambda xi: -4.0 * xi + 4.0*(1.0 - xi), \
                      lambda xi: 2.0 * xi + 2.0*(xi - 0.5)]
        self.__N = 3  #number of nodes in quadratic Lagrange polynomial

    def eval(self,n,xi):
        '''
        the function phi[n](xi), for any xi
        '''
        return self.__phi[n](xi)

    def ddx(self,n,xi):
        '''
        the function dphi[n](xi), for any xi
        '''
        return self.__dphi[n](xi)

    def size(self):
        '''
        the  number of points
        '''
        return self.__N

class FiniteElement(object):
    '''
    A single finite element
    FiniteElement(mesh,sfns,eltno,dofnos): constructor
        mesh is a Mesh
        sfns is a set of shape functions
        eltno=this element number, needs to be in mesh
        dofnos=numbers of this element's dofs (numDofs-sized array)
    endpts(): cell end points
    dofpts(): all dof locations
    dofnos(): all dof values
    numDofs(): number of dofs
    eval(n,x): phi[n](x)  (x, not xi)
    ddx(n,x):  dphi[n](x) (x, not xi)
    integral(f1=None,f2=None,derivative=False): integral(f1*f2*phi)
      f1, f2: ndof-sized vector of coeffs for local function
      derivative=True, do integral(f1*f2*dphi)
    '''

    def __init__(self,mesh,sfns,eltno,dofnos):
        '''
        mesh is the mesh it is built on
        sfns is the Shapefuns member
        eltno is this element's number
        endnos is a pair of ints giving the numbers of the endpoints
            in the mesh
        dofnos is an array of ints giving the numbers of the dofs
        '''
        # this element no. is same as mesh element no.
        assert(0 <= eltno < mesh.size())
        self.__eltno = eltno
#       print ( 'eltno=', eltno
        endnos = mesh.cells(eltno)
        assert(len(endnos) == 2)
        self.__endpts = np.array(mesh.coordinates(endnos))
        self.__numDofs = sfns.size()
        assert(sfns.size() == len(dofnos))
        self.__dofnos = dofnos
        self.__dofpts = np.linspace(self.__endpts[0],self.__endpts[1],self.__numDofs)
        self.__sfns = sfns
        #
        # Gauss points and weights: 3-pts are high enough for this
        #
        self.__gausspts = np.array(\
          (.112701665379258311482073460022,.5,.887298334620741688517926539978))
        self.__gausswts = np.array((5.0/18.0,8.0/18.0,5.0/18.0))
        # for efficiency, generate an array of shape functions evaluated
        # at the Gauss points
        self.__gaussvals = np.empty([self.__numDofs,self.__gausspts.size])
        for n in range(self.__numDofs):
            self.__gaussvals[n,:] = sfns.eval(n,self.__gausspts[:])

    def endpts(self):
        ''' access endpoints '''
        return self.__endpts

    def dofpts(self):
        ''' access dofpoints '''
        return self.__dofpts

    def dofnos(self):
        ''' access dof point numbers '''
        return self.__dofnos

    def numDofs(self):
        ''' access numDofs '''
        return self.__numDofs

    def eval(self,n,x):
        ''' 
        evaluate the n-th shape function on this element 
        at the spatial coordinate x
        '''
        # map x to xi
        xx = np.array(x)
        xi = (xx-self.__endpts[0])/(self.__endpts[1]-self.__endpts[0])
        # evaluate
        return self.__sfns.eval(n,xi) * (xi >= 0.)*(xi <= 1.)

    def ddx(self,n,x):
        ''' 
        evaluate the n-th shape function on this element 
        at the spatial coordinate x
        '''
        # map x to xi
        xi = (x-self.__endpts[0])/(self.__endpts[1]-self.__endpts[0])
        # evaluate
        return self.__sfns.ddx(n,xi)*(xi >= 0.)*(xi <= 1.0)

    def integral(self,f1=None,f2=None,derivative=False):
        '''
        Integrate either phi[i](xi)*f1(xi)*f2(xi) or dphi[i]*f1*f2
        over this element, depending on if derivative is False or True
        Returns a vector of 3 results, one for
        phi[0], one for phi[1], and one for phi[2].
        f1 and f2 are assumed to have been mapped to this element
          as arrays
        if derivative is True, phi is replaced with dphi
        '''
        L = self.__endpts[1]-self.__endpts[0]  # length of element
        t = self.__gausswts.copy()
        gp = self.__gausspts

        if not ( f1 is None ):
          assert(len(f1) == self.__numDofs)
          fvals = np.zeros([self.__gausspts.size])
          for n in range(self.__numDofs):
              fvals += f1[n]*self.__gaussvals[n,:]
          t *= fvals

        if not ( f2 is None ):
          assert(len(f2) == self.__numDofs)
          fvals = np.zeros([self.__gausspts.size])
          for n in range(self.__numDofs):
              fvals += f2[n]*self.__gaussvals[n,:]
          t *= fvals

        if derivative:
            # really: t *= L*(1/L)
            q = np.dot(np.array([self.__sfns.ddx(0,gp), \
                                 self.__sfns.ddx(1,gp), \
                                 self.__sfns.ddx(2,gp)]),t)
        else:
            t *= L  # correct for affine map x->xi
            q = np.dot(np.array([self.__sfns.eval(0,gp), \
                                 self.__sfns.eval(1,gp), \
                                 self.__sfns.eval(2,gp)]),t)

        return q

class FunctionSpace(object):
    '''
    A FunctionSpace has a list of elements
    numbered and with coords according to mesh
    FunctionSpace(mesh,sfns): constructor, sfns is ShapeFuns
    size(): number of elements
    ndofs(): number of all dofs
    dofpts(n=None): coordinates of dof[n] or all dofs
    int_phi_phi(c=None,derivative=[False,False]):
        integral(c*phi*phi) or
        integral(c*dphi*phi) or
        integral(c*dphi*dphi) or
        integral(c*phi*dphi)
    int_phi(f=None,derivative=False):
        integral(f*phi) or
        integral(f*dphi)

    '''
    def __init__(self,mesh,sfns):
        '''
        mesh is the mesh
        sfns is the Shapefuns
        '''
        self.__size = mesh.size()
        # number the elements in same way as mesh
        self.__elts = list([])
        self.__dofpts = list([])
        self.__nDOFs = 0
        for n in range(self.__size):
            # ASSUMING only boundary points are number 0 and (self.__size)
            if n == 0:
                self.__nDOFs += 3
                dofs = [2*n, 2*n+1, 2*n+2]
                newdofs = range(3)
            else:
                self.__nDOFs += 2
                dofs = [2*n, 2*n+1, 2*n+2]
                newdofs = range(1,3)
            fe = FiniteElement(mesh,sfns,n,dofs)
            self.__elts.append(fe)
            for i in newdofs:
                self.__dofpts.append(fe.dofpts()[i])
        self.__dofpts = np.array(self.__dofpts)

    def size(self):
        return len(self.__elts)

    def Ndofs(self):
        return self.__nDOFs

    def dofpts(self,n=None):
        if ( n is None ):
            return self.__dofpts
        else:
            return self.__dofpts[n]

    def int_phi_phi(self, c=None, derivative=[False,False]):
        '''
        assemble $\int c(x)\phi(x)\phi(x) dx$ or with $d\phi/dx$
        '''
        A = np.zeros([self.__nDOFs,self.__nDOFs])
        # loop over elements
        for elt in self.__elts:
            d0 = elt.dofnos()
            if not ( c is None ):
                cc = c[d0]
            else:
                cc = None
            N = elt.numDofs()
            endpts = elt.endpts()
            L = endpts[1]-endpts[0] # length of elt
            for j in range(N):
               if derivative[1]:
                   # chain rule: d(xi)/d(x) = 1/L
                   phi = elt.ddx(j,elt.dofpts())/L
               else:
                   phi = elt.eval(j,elt.dofpts())
               A[d0,d0[j]] += elt.integral(phi, cc, derivative=derivative[0])
        return A
            
    def int_phi(self, f=None, derivative=False):
        '''
        assemble $\int f(x)\phi(x) dx$ or with $d\phi/dx$
        '''
        F = np.zeros ( self.__nDOFs )

        for elt in self.__elts:
          d0 = elt.dofnos()
          if not ( f is None ):
            ff = f[d0]
          else:
            ff = None
          F[d0] += elt.integral ( ff,f2=None,derivative=derivative)

        return F
            
if __name__ == '__main__':
    ''' One test case '''

    N = 5
    rightpt = 5.0

    print ( '' )
    print ( 'FEM1D_CLASSES_TEST' )
    print ( '  Built-in test case with dx = ', rightpt / N )

    mesh = Mesh(N,0.0,rightpt)

    coords = mesh.coordinates()

    if N==5:
        print ( 'mesh.coordinates()=',coords )
        print ( 'mesh.cells()=', mesh.cells() )

    # generate an element

    sfns = Shapefns()
    print ( 'sfns.size()-3=',sfns.size()-3 )
    xi = np.linspace(0,1,100)

    import matplotlib.pyplot as plt
    if False:
        for n in range(3):
            plt.plot(xi,sfns.eval(n,xi))
            plt.show()
            plt.plot(xi,sfns.ddx(n,xi))
            plt.show()

    elt = FiniteElement(mesh,sfns,0,[0,1,2])

    if N==5 and np.abs(coords[-1]-5.0) < 1.e-10:
        # test some integrals
        print ( '' )
        print ( 'elt integral() err=', max(abs(elt.integral()-[1./6,2./3,1./6])) )
        print ( 'integral(deriv) err=', max(abs(elt.integral(derivative=True)- [-1,0,1])) )

        # test some more integrals
        # pick the function f(x)=x, find its expansion coefs
        ex = np.empty([sfns.size()])
        ex[0] = elt.endpts()[0]
        ex[2] = elt.endpts()[1]
        ex[1] = .5*(ex[0]+ex[2])
        ex2 = ex**2
        print ( 'integral(x) err=',max(abs(elt.integral(ex)-[0,1./3,1./6])) )
        print ( 'integral(x**2) err=',max(abs(elt.integral(ex2)-[-1./60,1./5,3./20])) )
        print ( 'integral(x**2) err=',max(abs(elt.integral(ex,ex)-[-1./60,1./5,3./20])) )

        print ( 'integral(x,phi) err=',\
              max(abs(elt.integral(ex,derivative=True)-[-1./6,-2./3,5./6])) )
        print ( 'integral(x**2,phi) err=',\
              max(abs(elt.integral(ex2,derivative=True)-[0,-2./3,2./3])) )
        print ( '' )

    V = FunctionSpace(mesh,sfns)
    print ( '' )
    print ( 'V.Ndofs()-correct=',V.Ndofs()-(2*N+1) )
    print ( 'V.size()-correct=',V.size()-N )

    x = V.dofpts()
    f = x.copy()
    print ( '' )
    print ( 'error in integral x over [',x[0],',',x[-1],']=',\
        np.sum(V.int_phi(f))-x[-1]**2/2. )
    f = 0.0*x+1
    print ( 'error in integral 1 over [',x[0],',',x[-1],']=',\
        np.sum(V.int_phi(f))-x[-1] )
    f = x.copy()**2
    print ( 'error in integral x**2 over [',x[0],',',x[-1],']=',\
        np.sum(V.int_phi(f))-x[-1]**3/3. )
    f = x.copy()**3
    print ( 'error in integral x**3 over [',x[0],',',x[-1],']=',\
        np.sum(V.int_phi(f))-x[-1]**4/4. )
    f = x.copy()**4
    print ( 'error in integral x**4 over [',x[0],',',x[-1],']=',\
        np.sum(V.int_phi(f))-x[-1]**5/5.,' (should be nonzero.)' )

    print ( '' )
    print ( 'norm(V.dofpts()-correct)=',\
        la.norm(V.dofpts()-np.linspace(0,coords[-1],2*N+1)) )

    # check eq \int \phi \phi * u = 1 gives 1 back (homog Neumann b.c.)
    # generate matrix by assembling $A_{ij}=\int\phi_i\phi_j$
    A = V.int_phi_phi()
    if N ==5 and np.abs(coords[-1]-5.0) < 1.e-10:
        print ( 'error A00=',A[0,0]-2./15. )
        print ( 'error A01=',A[0,1]-1./15. )
        print ( 'error A02=',A[0,2]+1./30. )
        print ( 'error A11=',A[1,1]-8./15. )
        print ( 'error A12=',A[1,2]-1./15. )
        print ( '' )
        print ( 'error A22=',A[2,2]-4./15. )
        print ( 'error A23=',A[2,3]-1./15. )
        print ( 'error A24=',A[2,4]+1./30. )
        print ( 'error A33=',A[3,3]-8./15. )
        print ( 'error A34=',A[3,4]-1./15. )
        print ( '' )

    Ndofs = V.Ndofs()
    f = np.random.rand(Ndofs)
    b = V.int_phi(f)
    print ( 'norm(A*f-b)=',la.norm(np.dot(A.transpose(),f)-b) )

    #trivial check with coefficient
    c = np.ones([Ndofs])
    A1 = V.int_phi_phi(c)
    print ( 'Norm difference matrices=',la.norm(A-A1) )

    # try putting a coefficient in: c(x) = (1+x)
    # rhs = ones, soln = 1/(1+x)
    c = (1.0+x)
    B = V.int_phi_phi(c)
    if N == 5 and np.abs(coords[-1]-5.0) < 1.e-10:
        print ( '' )
        print ( 'error B00=',B[0,0]-3./20. )
        print ( 'error B01=',B[0,1]-1./15. )
        print ( 'error B02=',B[0,2]+1./20. )
        print ( 'error B11=',B[1,1]-12./15. )
        print ( 'error B12=',B[1,2]-2./15. )
        print ( '' )
        print ( 'error B22=',B[2,2]-8./15. )
        print ( 'error B23=',B[2,3]-2./15. )
        print ( 'error B24=',B[2,4]+1./12. )
        print ( 'error B33=',B[3,3]-4./3. )
        print ( 'error B34=',B[3,4]-3./15. )

    C = V.int_phi_phi(derivative=[True,True])
    if N == 5 and np.abs(coords[-1]-5.0) < 1.e-10:
        print ( '' )
        print ( 'Laplace Matrix' )
        print ( 'error C00*3=',C[0,0]-7./3. )
        print ( 'error C01*3=',C[0,1]+8./3. )
        print ( 'error C02*3=',C[0,2]-1./3. )
        print ( 'error C11*3=',C[1,1]-16./3. )
        print ( 'error C12*3=',C[1,2]+8./3. )
        print ( '' )
        print ( 'error C22*3=',C[2,2]-14./3. )
        print ( 'error C23*3=',C[2,3]+8./3. )
        print ( 'error C24*3=',C[2,4]-1./3. )
        print ( 'error C33*3=',C[3,3]-16./3. )
        print ( 'error C34*3=',C[3,4]+8./3. )
        print ( '' )

    soln2 = np.ones([Ndofs])
    b2 = np.dot(C,soln2)
    print ( 'const soln Laplace, norm check=',la.norm(b2) )

    soln = x
    b0 = np.dot(C,soln)
    rhs0 = V.int_phi(np.zeros([Ndofs]))
    # natural b.c. not satisfied, don't check them
    rhs0[0] = -b0[0]
    rhs0[-1] = -b0[-1]
    print ( 'soln=x Laplace, norm check=',la.norm(rhs0+b0) )

    soln = x**2
    b1 = np.dot(C,soln)
    rhs1 = V.int_phi(2.0*np.ones([Ndofs]))
    # natural b.c. not satisfied on right, don't check it 
    rhs1[-1] = -b1[-1]
    print ( 'soln=x**2 Laplace, norm check=',la.norm(rhs1+b1) )

    D = V.int_phi_phi(derivative=[False,True])
    soln = np.ones([V.Ndofs()])
    b2 = np.dot(D,soln)
    print ( 'norm check (rhs d/dx+Neumann, const soln)=',la.norm(b2) )

    D[0,0] = 1.0
    D[0,1:] = 0.0
    D[-1,-1] = 1.0
    D[-1,0:-1] = 0.0
    soln = x
    b3 = np.dot(D,soln)
    rhs3 = V.int_phi(np.ones([Ndofs]))
    rhs3[0] = soln[0]
    rhs3[-1] = soln[-1]
    print ( 'norm check (d/dx+Dirichlet soln=x)=',la.norm(rhs3-b3) )

    print ( '' )
    print ( 'FEM1D_CLASSES' )
    print ( '  Normal end of execution.' )