module fem
use precision
implicit none

integer,private :: ndof=-1  !! -1 is flag for uninitialized
real(PREC),allocatable :: Mat(:,:),Mat1(:,:),Mat2(:,:),Mat3(:,:),rhs(:)
real(PREC),allocatable :: mesh(:)
real(PREC),private :: dx
real(PREC),private,parameter :: zero=0._PREC,one=1._PREC

contains

subroutine makemat(nin)
  integer,intent(in) :: nin
  integer :: n

  if (allocated(mesh)) stop 'fem: already allocated!'
  if (ndof>0) stop 'fem: n already defined!'

  ndof=nin
  allocate(mesh(ndof),Mat(ndof,ndof),Mat1(ndof,ndof),Mat2(ndof,ndof),Mat3(ndof,ndof))

  dx=1.0_PREC/(ndof+1)

  mesh=(/ (dx*n,n=1,ndof) /)

  !! generate matrices
  do n=1,ndof
    Mat1(n,n)=-gaussquad(phip,n,phip,n,max(zero,(n-2)*dx),min(one,(n+2)*dx));
    Mat2(n,n)=gaussquad(phi,n,phip,n,max(zero,(n-2)*dx),min(one,(n+2)*dx));
    Mat3(n,n)=gaussquad(phi,n,phi,n,max(zero,(n-2)*dx),min(one,(n+2)*dx));
    if (n<ndof) then
      Mat1(n,n+1)=-gaussquad(phip,n,phip,n+1,max(zero,(n-2)*dx),min(one,(n+2)*dx));
      Mat1(n+1,n)=-gaussquad(phip,n,phip,n+1,max(zero,(n-2)*dx),min(one,(n+2)*dx));
      Mat2(n,n+1)=gaussquad(phi,n,phip,n+1,max(zero,(n-2)*dx),min(one,(n+2)*dx));
      Mat2(n+1,n)=gaussquad(phi,n+1,phip,n,max(zero,(n-2)*dx),min(one,(n+2)*dx));
      Mat3(n,n+1)=gaussquad(phi,n,phi,n+1,max(zero,(n-2)*dx),min(one,(n+2)*dx));
      Mat3(n+1,n)=gaussquad(phi,n,phi,n+1,max(zero,(n-2)*dx),min(one,(n+2)*dx));
    endif
    if (n<ndof-1) then
      Mat1(n+2,n)=-gaussquad(phip,n,phip,n+2,max(zero,(n-2)*dx),min(one,(n+2)*dx));
      Mat1(n,n+2)=-gaussquad(phip,n,phip,n+2,max(zero,(n-2)*dx),min(one,(n+2)*dx));
      Mat2(n+2,n)=gaussquad(phi,n+2,phip,n,max(zero,(n-2)*dx),min(one,(n+2)*dx));
      Mat2(n,n+2)=gaussquad(phi,n,phip,n+2,max(zero,(n-2)*dx),min(one,(n+2)*dx));
      Mat3(n+2,n)=gaussquad(phi,n,phi,n+2,max(zero,(n-2)*dx),min(one,(n+2)*dx));
      Mat3(n,n+2)=gaussquad(phi,n,phi,n+2,max(zero,(n-2)*dx),min(one,(n+2)*dx));
    endif
  enddo
  Mat=Mat1+Mat2+Mat3;

end subroutine makemat

subroutine makerhs(frhs)
  interface
    function frhs(k,x)
      use precision
      integer,intent(in)    :: k
      real(PREC),intent(in) :: x
      real(PREC)            :: frhs
    end function frhs
  end interface
  integer m

  allocate(rhs(ndof))
  do m=1,ndof
    rhs(m)=gaussquad(frhs,m,phi,m,max(0._PREC,(m-2)*dx),min(1._PREC,(m+2)*dx));
  enddo

end subroutine makerhs

function exact9(x)
  ! exact solution to y''+y'+y=x+1
  real(PREC) :: x(:)
  real(PREC),allocatable :: exact9(:)
  real(PREC) :: omega
  real(PREC),parameter :: one=1.0_PREC,two=2.0_PREC,three=3.0_PREC

  allocate(exact9(size(x)))

  omega=sqrt(three)/two;
  exact9=x-exp(-(x-one)/two)*sin(omega*x)/sin(omega)
end function exact9

function phi(n,x)
  ! Lagrange quadratic basis functions
  integer,intent(in)    :: n
  real(PREC)            :: phi
  real(PREC),intent(in) :: x
  real(PREC),parameter  :: two=2.0_PREC
  
  if(mod(n,2)==0) then  ! n is even
  
    if ((n-2)*dx <= x .and. x <= n*dx) then
      phi= (x-(n-1)*dx)*(x-(n-2)*dx)/(two*dx**2);
    else if (n*dx <= x .and. x <= (n+2)*dx) then
      phi= ((n+1)*dx-x)*((n+2)*dx-x)/(two*dx**2);
    else
      phi=0.0_PREC;
    endif
    
  else  ! n is odd
  
    if ((n-1)*dx <= x .and. x <= (n+1)*dx) then
      phi= ((n+1)*dx-x)*(x-(n-1)*dx)/(dx**2);
    else
      phi=0.0_PREC;
    endif
  
  endif
end function phi

function phip(n,x)
  ! Lagrange quadratic basis functions
  integer,intent(in)    :: n
  real(PREC)            :: phip
  real(PREC),intent(in) :: x
  real(PREC),parameter  :: two=2.0_PREC
  
  if(mod(n,2)==0) then  ! n is even
  
    if ((n-2)*dx <= x .and. x <= n*dx) then
      phip = (one)*(x-(n-2)*dx)/(two*dx**2) &
        &  + (x-(n-1)*dx)*(one)/(two*dx**2);
    else if (n*dx <= x .and. x <= (n+2)*dx) then
      phip= ((n+1)*dx-x)*(-one)/(two*dx**2) &
        & + (-one)*((n+2)*dx-x)/(two*dx**2);
    else
      phip=0.0_PREC;
    endif
    
  else  ! n is odd
  
    if ((n-1)*dx <= x .and. x <= (n+1)*dx) then
      phip= (-one)*(x-(n-1)*dx)/(dx**2) &
        & + ((n+1)*dx-x)*(one)/(dx**2);
    else
      phip=0;
    endif
  
  endif
end function phip

function gaussquad(f1,k1,f2,k2,a,b)
  ! compute the integral from a to b of f1(k1,x)*f2(k2,x)dx
  ! using 3-point Gaussian quadrature
  ! f1 and f2 are functions
  ! The functions f1 and f2 are assumed to be piecewise polynomials
  !   with pieces [a+n*dx,a+(n+1)*dx]
  ! If the product f1*f2 is no higher than fifth degree, the
  ! 3-point Gauss formula used here is exact.

  integer,intent(in) :: k1,k2
  real(PREC),intent(in) :: a,b
  real(PREC) :: quadpt,wa,wb,wc,gaussquad
  real(PREC) :: xleft,xright,len,xa,xb,xc
  integer    :: count

  interface
    function f1(k,x)
      use precision
      integer,intent(in) :: k
      real(PREC),intent(in) :: x
      real(PREC) :: f1
    end function f1
    function f2(k,x)
      use precision
      integer,intent(in) :: k
      real(PREC),intent(in) :: x
      real(PREC) :: f2
    end function f2
  end interface
  
  ! parameters for 3-point Gauss quadrature over [-1,1]
  ! quantities are halved because length of [-1,1] is 2
  
  quadpt=0.774596669241483377035853079956_PREC/2.0_PREC;
  wa=5.0_PREC/18.0_PREC;
  wb=4.0_PREC/9.0_PREC;
  wc=5.0_PREC/18.0_PREC;
  
  if (b<=a) stop 'gaussquad: interval [a,b] must have positive length'
  if (dx<=0) stop 'gaussquad: dx must be >0'
  
  xleft=a;
  xright=min(a+dx,b);
  len=xright-xleft;
  

  gaussquad=0;
  do
    if (len<= 100.0_PREC*epsilon(dx)*dx) exit
  
    xb=(xleft+xright)/2.0_PREC;
    xa=xb-len*quadpt;
    xc=xb+len*quadpt;
  
    gaussquad=gaussquad+ len*(wa*f1(k1,xa)*f2(k2,xa) + &
             &  wb*f1(k1,xb)*f2(k2,xb) + &
             &  wc*f1(k1,xc)*f2(k2,xc) );
  
    xleft=xright;
    xright=min(xleft+dx,b);
    len=xright-xleft;
  
  enddo
end function gaussquad

end module fem