#! /usr/bin/env python3
#
def phi ( i, xn, xs ):

#*****************************************************************************80
#
## phi() evaluates piecewise quadratic Lagrange basis function I on a mesh XN.
#
  import numpy as np
#
#  We need to ensure that xs can be treated as a vector.
#
  xs = np.atleast_1d ( xs )
#
#  Set up zs, which will contain the function value at each point xs.
#
  ns = len ( xs )
  zs = np.zeros ( ns )
#
#  Formula for odd index I:
#
  if ( ( i % 2 ) == 1 ):

    ks = np.nonzero ( ( xn[i-1] <= xs ) & ( xs <= xn[i+1] ) )
    zs[ks] =   ( xs[ks] - xn[i-1] ) / ( xn[i] - xn[i-1] ) \
             * ( xs[ks] - xn[i+1] ) / ( xn[i] - xn[i+1] )
#
#  Formula for even index I:
#
  else:

    nn = len ( xn )

    if ( 0 < i ):
      ks = np.nonzero ( ( xn[i-2] <= xs ) & ( xs <= xn[i] ) )
      zs[ks] = ( xs[ks] - xn[i-2] ) / ( xn[i] - xn[i-2] ) \
             * ( xs[ks] - xn[i-1] ) / ( xn[i] - xn[i-1] )

    if ( i < nn - 1 ):
      ks = np.nonzero ( ( xn[i] <= xs ) & ( xs <= xn[i+2] ) )
      zs[ks] = ( xs[ks] - xn[i+2] ) / ( xn[i] - xn[i+2] ) \
             * ( xs[ks] - xn[i+1] ) / ( xn[i] - xn[i+1] )

  return zs

def phi_test ( ):

#*****************************************************************************80
#
## phi_test() tests phi().
#
  import matplotlib.pyplot as plt
  import numpy as np

  print ( "phi_test():" )

  a = 2.0
  b = 8.0
  nn = 7
  xn = np.linspace ( a, b, nn )

  ns = 101
  xs = np.linspace ( a, b, ns )
  ys = np.zeros ( ns )

  i = 3
  ys = phi ( i, xn, xs )

  plt.plot ( xs, ys )
  plt.grid ( True )
  plt.title ( "Phi(" + str ( i ) + ")(x)" )
  plt.savefig ( 'phi.png' )
  plt.show ( )
  plt.close ( )

def phip ( i, xn, xs ):

#*****************************************************************************80
#
## phip(): derivative, piecewise quadratic Lagrange basis function I on mesh XN.
#
  import numpy as np
#
#  We need to ensure that xs can be treated as a vector.
#
  xs = np.atleast_1d ( xs )
#
#  Set up zps, which will contain the derivative value at each point xs.
#
  ns = len ( xs )
  zps = np.zeros ( ns )
#
#  Formula for odd index I:
#
  if ( ( i % 2 ) == 1 ):

    ks = np.nonzero ( ( xn[i-1] <= xs ) & ( xs <= xn[i+1] ) )
    zps[ks] =     1.0                / ( xn[i] - xn[i-1] ) * ( xs[ks] - xn[i+1] ) / ( xn[i] - xn[i+1] ) \
              + ( xs[ks] - xn[i-1] ) / ( xn[i] - xn[i-1] ) * 1.0                  / ( xn[i] - xn[i+1] )
#
#  Formula for even index I:
#
  else:

    nn = len ( xn )

    if ( 0 < i ):
      ks = np.nonzero ( ( xn[i-2] <= xs ) & ( xs <= xn[i] ) )
      zps[ks] =   1.0                / ( xn[i] - xn[i-2] ) * ( xs[ks] - xn[i-1] ) / ( xn[i] - xn[i-1] ) \
              + ( xs[ks] - xn[i-2] ) / ( xn[i] - xn[i-2] ) *   1.0                / ( xn[i] - xn[i-1] )

    if ( i < nn - 1 ):
      ks = np.nonzero ( ( xn[i] <= xs ) & ( xs <= xn[i+2] ) )
      zps[ks] =   1.0                / ( xn[i] - xn[i+2] ) * ( xs[ks] - xn[i+1] ) / ( xn[i] - xn[i+1] ) \
              + ( xs[ks] - xn[i+2] ) / ( xn[i] - xn[i+2] ) *   1.0                / ( xn[i] - xn[i+1] )

  return zps

def phip_test ( ):

#*****************************************************************************80
#
## phip_test() tests phip().
#
  import matplotlib.pyplot as plt
  import numpy as np

  print ( "phip_test:" )
  a = 2.0
  b = 8.0
  nn = 7
  xn = np.linspace ( a, b, nn )

  ns = 101
  xs = np.linspace ( a, b, ns )

  i = 4
  ysp = phip ( i, xn, xs )

  plt.plot ( xs, ysp )
  plt.grid ( True )
  plt.title ( "Phi'(" + str ( i ) + ")(x)" )
  plt.savefig ( 'phip.png' )
  plt.show ( )
  plt.close ( )

def lagrange_basis_test ( ):

#*****************************************************************************80
#
## lagrange_basis_test() tests lagrange_basis().
#
  print ( "lagrange_basis_test" )

  phi_test ( )
  phip_test ( )

if ( __name__ == '__main__' ):
  lagrange_basis_test ( )