"""
0. Solve Laplace eqn using FENICs
1. Evaluate a given objective function
2. Compute an approximation to directional derivative using finite differences
3. Evaluate the directional derivative using the "direct" method
"""

from dolfin import *

N = 64  # number of elements in the mesh
mesh = UnitSquareMesh(N,N)

# state space
Y = FunctionSpace(mesh,"Lagrange",1)
y_ = TrialFunction(Y)
v_ = TestFunction(Y)
# control space
U = FunctionSpace(mesh,"Discontinuous Lagrange",0)
u  = Function(U)

# define the bilinear form
a = inner(grad(y_),grad(v_))*dx
# define the right hand side
# let beta be the indicator function for a circle in the middle
beta = Expression("(x[0]-0.5)*(x[0]-0.5)+(x[1]-0.5)*(x[1]-0.5)<=0.5*0.5", degree=0)
l = inner(beta*u,v_)*dx

# define the Dirichlet boundary conditions
class DirichletBoundary(SubDomain):
    def inside(self, x, on_boundary):
        # in this example we mark all boundary as Dirichlet
        return on_boundary

y0 = Constant(1.0)
bc = DirichletBC(Y,y0,DirichletBoundary())


# assemble the system and apply boundary conditions
A = assemble(a)
# before we assemble the right hand side, we need to assign some value to our control variable
project(Constant(1.0),U,function=u)
b = assemble(l)
bc.apply(A,b)

# solve the system
y = Function(Y)
solve(A,y.vector(),b)

# define the objective functional
lam = 1.0
y_Omega = Expression("x[0]*x[0] + x[1]*x[1]", degree=2)
J = 0.5*inner(y-y_Omega,y-y_Omega)*dx + lam/2*inner(u,u)*dx


# evaluate the objective functional at a present control point
J0 = assemble(J)

# perturb the solution a little bit in some random direction h
from numpy import random, arange
h = Function(U)
h.vector().set_local(0.5-random.random(h.vector().size()))


"""
We will now compute the directional derivative of J using a different method.
J = 0.5 ||y-y_Omega||^2 + lam/2 ||u||^2,
and y = y0 + Su, where S is the solution operator for the state problem with homogeneous boundary conditions.

Thus we have: 
J'h = J'_y Sh + J'_u h = 
    = (y-y_Omega, Sh) + lam(u,h)
"""

# first, we homogenize the boundary conditions
bc0 = DirichletBC(bc)
bc0.homogenize()
# second, we change the right-hand side of the problem to "beta*h"
A = assemble(a)
b = assemble(beta*h*v_*dx)
bc0.apply(A,b)
# put dy = Sh
dy = Function(Y)
solve(A,dy.vector(),b)
# evaluate the derivative now
dJdh = assemble(inner(y-y_Omega,dy)*dx + lam*inner(u,h)*dx)
print("dJdh = %e" % (dJdh))


# save the old u
u_old = Function(U)
u_old.vector().set_local(u.vector().array())

# run through a sequence of perturbations
for delta in pow(10.,-arange(10)):
    u.vector().set_local(u_old.vector().array() + delta*h.vector().array())
    # reassemble the right hand side and resolve the system
    A = assemble(a)
    b = assemble(l)
    bc.apply(A,b)
    solve(A,y.vector(),b)
    # re-evaluate the functional
    J1 = assemble(J)
    # compute the FD approximation
    dJdh = (J1-J0)/delta
    print("delta = %e, dJdh = %e" % (delta,dJdh))