"""
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
4. Evaluate the directional derivative using the adjoint 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(inner(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(direct) = %e" % (dJdh))

"""
We use adjoint method now.

First, adjoint problem is

-div(grad(p)) = J'_y = (y-y_Omega,.)
p_dOmega = 0

and then
J'h = S*(J'_y)h + J'_u h = 
    = (beta*p,h) + lam(u,h) 
"""

# Adjoint problem: same left hand side 
A = assemble(a)
# but we change the right hand side to "y-y_Omega"
b = assemble(inner(y-y_Omega,v_)*dx)
bc0.apply(A,b)
# p is going to be our adjoint state
p = Function(Y)
solve(A,p.vector(),b)
# evaluate the derivative now
dJdh = assemble(inner(beta*p+lam*u,h)*dx)
print("dJdh(adjoint) = %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))