#! /usr/bin/env python
#
"""
I copied the code from fem1d_model.py and have made some adjustments
"""
import numpy as np
import scipy.linalg as la

def fem1d_model ( n ):
#
## FEM1D_MODEL solves a 1D "model" boundary value problem using finite elements.
#
#  Location:
#
#    http://people.sc.fsu.edu/~jburkardt/py_src/fem1d/fem1d_model.py
#
#  Discussion:
#
#    The PDE is defined for 0 < x < 1:
#      -u'' + u = x
#    with boundary conditions
#      u(0) = 0,
#      u(1) = 0.
#
#    The exact solution is:
#      exact(x) = x - sinh(x) / sinh(1.0)
#
#    This program is different from FEM1D.PY:
#    * the problem to be solved is different, and includes a linear term;
#    * the code to assemble the matrix is different.  We evaluate all the
#      basis functions and derivatives, and then form the combinations
#      that must be added to the system matrix and right hand side.
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license. 
#
#  Modified:
#
#    23 September 2014
#
#  Author:
#
#    John Burkardt
#
#  Local parameters:
#
#    Local, integer N, the number of elements.
#

#
#  The mesh will use N+1 points between A and B.
#  These will be indexed X[0] through X[N].
#
  a = 0.0
  b = 1.0
  x = np.linspace ( a, b, n + 1 )

    
#
#  Set a 3 point quadrature rule on the reference interval [0,1].
#
  ng = 3

  xg = np.array ( ( \
    0.112701665379258311482073460022, \
    0.5, \
    0.887298334620741688517926539978 ) )

  wg = np.array ( ( \
    5.0 / 18.0, \
    8.0 / 18.0, \
    5.0 / 18.0 ) )
#
#  Compute the system matrix A and right hand side RHS.
#
  A = np.zeros ( ( n + 1, n + 1 ) )
  rhs = np.zeros ( n + 1 )
#
#  Look at element E: (0, 1, 2, ..., N-1).
#
  for e in range ( 0, n ):

    l = e
    r = e + 1

    xl = x[l]
    xr = x[r]
#
#  Consider quadrature point Q: (0, 1, 2 ) in element E.
#
    for q in range ( 0, ng ):
#
#  Map XG and WG from [0,1] to
#      XQ and QQ in [XL,XR].
#
      xq = xl + xg[q] * ( xr - xl )
      wq = wg[q] * ( xr - xl )
#
#  Evaluate at XQ the basis functions and derivatives for XL and XR.
#
      phil = ( xr - xq  ) / ( xr - xl )
      philp = - 1.0 / ( xr - xl )

      phir = ( xq - xl ) / ( xr - xl )
      phirp = 1.0 / ( xr - xl )
#
#  Compute the following contributions:
#
#    L,L  L,R  L,Fx
#    R,L  R,R  R,Fx
#
      A[l][l] = A[l][l] + wq * ( philp * philp + phil * phil )
      A[l][r] = A[l][r] + wq * ( philp * phirp + phil * phir )
      rhs[l]  = rhs[l]  + wq *                   phil * rhs_fn ( xq )

      A[r][l] = A[r][l] + wq * ( phirp * philp + phir * phil )
      A[r][r] = A[r][r] + wq * ( phirp * phirp + phir * phir )
      rhs[r]  = rhs[r]  + wq *                   phir * rhs_fn ( xq )
#
#  Modify the linear system to enforce the left boundary condition.
#
  A[0,0] = 1.0
  A[0,1:n+1] = 0.0
  rhs[0] = 0.0
#
#  Modify the linear system to enforce the right boundary condition.
#
  A[n,n] = 1.0
  A[n,0:n] = 0.0
  rhs[n] = 0.0
#
#  Solve the linear system.
#
  u = la.solve ( A, rhs )

  l2_error,h1_error = compute_error(u,exact_fn,exact_fn_der,x,xg,wg) 
  u_ave  = compute_average(u,x,xg,wg)
  
  return l2_error, h1_error, u_ave
  

def exact_fn ( x ):
#
## EXACT_FN evaluates the exact solution.
#
  from math import sinh
  value = x - sinh ( x ) / sinh ( 1.0 )
  return value

def exact_fn_der( x ):
  from math import cosh,sinh
  # Evaluate the derivative of the exact function
  return 1 - cosh( x ) / sinh( 1.0)
  
def rhs_fn ( x ):
#
## RHS_FN evaluates the right hand side.
#
  value = x
  return value

# ---------------------------------------------------------------------------
# Write a subroutine to compute the integral of a function 
# ---------------------------------------------------------------------------
def compute_error(ua,ue,uep,x,xg,wg):
  """ 
  Compute the integral of a function, using a composite Gauss rule over
  all elements. 
  Inputs: 
    ua       - double, array of finite element coefficients
    ue       - function, exact solution
    uep      - function, derivative of the exact solution
    x        - double, list of finite element nodes. 
    xg       - double, Gauss nodes on the interval [0,1]
    wg       - double, Gauss weights
    
  Outputs: 
    error_l2 - double, L2 error
    error_h1 - double, H1 error
  """
  # Initialize the the L2 and H10 errors
  l2_error = 0.0
  h1_error = 0.0
  
  n = len(x)-1
  
  for e in range(n):
    
    l = e
    r = e + 1

    xl = x[l]
    xr = x[r]
  
    # Consider quadrature point Q: (0, 1, 2 ) in element E.
    for q in range ( len(xg) ):
      #  Map XG and WG from [0,1] to
      #      XQ and QQ in [XL,XR].
      xq = xl + xg[q] * ( xr - xl )
      wq = wg[q] * ( xr - xl )
      
      #  Evaluate at XQ the basis functions and derivatives for XL and XR.
      phil = ( xr - xq  ) / ( xr - xl )
      philp = - 1.0 / ( xr - xl ) 

      phir = ( xq - xl ) / ( xr - xl )
      phirp = 1.0 / ( xr - xl )


      # Evaluate at XQ the finite element approximation
      ua_loc  = ua[l]*phil + ua[r]*phir
      uap_loc = ua[l]*philp + ua[r]*phirp
  
      #  Evaluate the exact function and its derivative at XQ
      ue_loc  = exact_fn(xq)
      uep_loc = exact_fn_der(xq)
    
      # Compute the 
      l2_error = l2_error + wq * (ua_loc - ue_loc)**2.0
      h1_error = h1_error + wq * (uap_loc - uep_loc)**2.0
    

  # Take square roots
  l2_error = np.sqrt(l2_error)
  h1_error = np.sqrt(h1_error)

  return l2_error, h1_error

def compute_average(ua,x,xg,wg):
  """ 
  Compute the spatial average of a finite element function ua 
  
  Inputs:
    ua - double, finite element coefficients of the function
    x  - double, finite element nodes
    xg - double, Gaussian quadrature nodes on [0,1]
    wg - double, Gaussian quadrature weights on [0,1]
    
  Outputs:
    u_ave - double, spatial average of ua
  """
  # Initialize the average
  ua_ave = 0.0

  # number of elements
  n = len(x)-1
  
  for e in range(n):
    
    l = e
    r = e + 1

    xl = x[l]
    xr = x[r]
  
    # Consider quadrature point Q: (0, 1, 2 ) in element E.
    for q in range ( len(xg) ):
      #  Map XG and WG from [0,1] to
      #      XQ and QQ in [XL,XR].
      xq = xl + xg[q] * ( xr - xl )
      wq = wg[q] * ( xr - xl )
      
      #  Evaluate at XQ the basis functions for XL and XR.
      phil = ( xr - xq  ) / ( xr - xl )
      phir = ( xq - xl ) / ( xr - xl )

      # Evaluate at XQ the finite element approximation
      ua_loc  = ua[l]*phil + ua[r]*phir
   
      # Update the average
      ua_ave += wq * ua_loc
      
  return ua_ave
    
#
#  If this script is called directly, then run it as a program.
#
if ( __name__ == '__main__' ):
  nlist  = [2,4,8,16,32,64]
  n_runs = len(nlist)
  l2_error = [];
  h1_error = [];
    
  print ""
  print "Question 1:"
  print "N    L2 ratios    H1 ratios" 
  for n in nlist:
    l2_err_temp,h1_err_temp,garbage = fem1d_model ( n )
    l2_error.append(l2_err_temp)
    h1_error.append(h1_err_temp)

  for i in range(len(nlist)-1):
    print "%2d   %f     %f" % (nlist[i+1], l2_error[i]/l2_error[i+1],h1_error[i]/h1_error[i+1])
  
  print ""
  print "Question 2:"
  print "The spatial average of u"
  print "n   u_ave"
  nlist = [2,4,8]
  u_ave = [];
  for n in nlist:
    garbage1,garbage2,u_ave_temp = fem1d_model(n)
    u_ave.append(u_ave_temp)
    print "%d   %f" % (n, u_ave_temp)