# -*- coding: utf-8 -*-
"""
Description: Investigate the accuracy of piecewise polynomial interpolation for 
             functions f and g
             
Created on Fri Sep 26 13:56:31 2014

@author: hans-werner
"""

from dolfin import *
import matplotlib.pyplot as plt
import numpy

# -----------------------------------------------------------------------------
# Problem a)
# -----------------------------------------------------------------------------
# Define the mesh
mesh = UnitIntervalMesh(5)

# Define f and g 
f = Expression("exp(-x[0]*x[0])*sin(4*pi*x[0])", cell = mesh.ufl_cell())
g = Expression("fabs(x[0]-pi/4)/(x[0]-pi/4)", cell = mesh.ufl_cell())

# Define the Function Spaces
V1 = FunctionSpace(mesh,"CG",1)
V2 = FunctionSpace(mesh,"CG",2)
V3 = FunctionSpace(mesh,"CG",3)

# Define the interpolants
# f
fi1 = interpolate(f,V1)
fi2 = interpolate(f,V2)
fi3 = interpolate(f,V3)
# g
gi1 = interpolate(g,V1)
gi2 = interpolate(g,V2)
gi3 = interpolate(g,V3)

# Print out the dimension of V1,V2, and V3
print "Dimension of V1 %d" % (V1.dim()) 
print "Dimension of V2 %d" % (V2.dim())
print "Dimension of V3 %d" % (V3.dim())


# -----------------------------------------------------------------------------
# Problem b) 
# -----------------------------------------------------------------------------

def evaluate_fn(f,x):
  """ Evaluate a function at all points in a vector x
  Inputs:
  f - function
  x - numpy array
  
  Outputs:
  fv = f(xi) for all xi in x 
  """
  n_pts = x.size
  fv = numpy.zeros([n_pts,1])
  for i in range(n_pts):
    fv[i] = f(x[i])
    
  return fv
  
 

# Plot f and its interpolants
x = numpy.linspace(0.0,1.0,51)
f_v   = evaluate_fn(f,x)
fi1_v = evaluate_fn(fi1,x)
fi2_v = evaluate_fn(fi2,x)
fi3_v = evaluate_fn(fi3,x)

plt.plot(x,f_v,x,fi1_v,x,fi2_v,x,fi3_v)
plt.legend( ('f','fi1','fi2','fi3'))


# Plot g and its interpolants
g_v   = evaluate_fn(g,x)
gi1_v = evaluate_fn(gi1,x)
gi2_v = evaluate_fn(gi2,x)
gi3_v = evaluate_fn(gi3,x)
  
plt.figure()
plt.plot(x,g_v,x,gi1_v,x,gi2_v,x,gi3_v)
plt.legend( ('g','gi1','gi2','gi3'), loc = 0)

# -----------------------------------------------------------------------------
# Compute the maximum eror
# -----------------------------------------------------------------------------
def compute_max_error(f_exact,f_approx,x):
  """ Compute the maximum absolute distance between f_exact and f_approx,
      based on the points given in x.
  """
  fe_v = evaluate_fn(f_exact,x)
  fa_v = evaluate_fn(f_approx,x)
  
  max_error = np.max(np.abs(fe_v - fa_v))
  return max_error

# f :
ef1 = compute_max_error(f,fi1,x)
ef2 = compute_max_error(f,fi2,x)
ef3 = compute_max_error(f,fi3,x)

# g : 
eg1 = compute_max_error(g,gi1,x)
eg2 = compute_max_error(g,gi2,x)
eg3 = compute_max_error(g,gi3,x)

# Print
print "V1: error(f) = %f, error(g) = %f" % ( ef1,eg1 )
print "V2: error(f) = %f, error(g) = %f" % ( ef2,eg2 )
print "V3: error(f) = %f, error(g) = %f" % ( ef3,eg3 )   



# -----------------------------------------------------------------------------
# Problem d)
# -----------------------------------------------------------------------------
cell_num_list = [5,10,20,40]

for n in cell_num_list:
  print "Number of cells = %d" %(n) 
  mesh = UnitIntervalMesh(n)
  V    = FunctionSpace(mesh,"CG",1)
  
  fi   = interpolate(f,V)
  gi   = interpolate(g,V)
  
  error_fi = compute_max_error(f,fi,x)
  error_gi = compute_max_error(g,gi,x)
  
  print "|f - fi| = %f" % (error_fi)
  print "|g - gi| = %g" % (error_gi)
  
  
# Define the initial mesh and function space
mesh = IntervalMesh(10,0,1)
V    = FunctionSpace(mesh,"Lagrange",2)

# Define the functions to be interpolated
meshe = IntervalMesh(1000,0,1)
Ve    = FunctionSpace(meshe,"Lagrange",1)



# Plot the interpolant
fh = interpolate(f,V)
gh = interpolate(g,V)

x = numpy.linspace(0,1,100)
fh_val = numpy.zeros([100,1])
for i in range(x.size):
    fh_val[i] = fh(x[i])

fh_at_nodes = numpy.zeros([mesh.num_vertices(),1])    
for i in range(mesh.num_vertices()):
    fh_at_nodes[i] = fh(mesh.coordinates()[i])       
    
#plt.plot(x,fh_val,mesh.coordinates(),fh_at_nodes,'ro')



#fh_at_fine = interpolate(fh,Ve)

#plot(fh_at_fine,V)
#integral = assemble(fh_at_fine*dx)
#print(integral)