"""
Solve the unconstrained control problem:

minimize J = 0.5 ||y-y_Omega||^2 + lam/2 ||u||^2,
subject to  -div(grad(y)) = beta*u,  in Omega
            y = y0  on dOmega


the solution is given by the first order optimality conditions:

J' = (beta p + lam u,.) = 0

meaning u = -1/lam beta p,

where p is the solution to the adjoint problem:

-div(grad(p)) = y-y_Omega, in Omega
p = 0, on dOmega


Substituting gives us a couple of linear PDEs:

-div(grad(y)) + 1/lam beta^2 p = 0,
-div(grad(p)) - y              = -y_Omega,

which we can formulate as a variational problem and solve.
"""

from dolfin import *

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

# state space
Y = FunctionSpace(mesh,"Lagrange",1)
# control space
U = FunctionSpace(mesh,"Discontinuous Lagrange",0)
# In order to solve the coupled problem, we form a product space Y*Y:
Y2 = MixedFunctionSpace([Y,Y])
y_,p_ = TrialFunctions(Y2)
v_,q_ = TestFunctions(Y2)


lam = Constant(1.0E-06)
beta = Expression("(x[0]-0.5)*(x[0]-0.5)+(x[1]-0.5)*(x[1]-0.5)<=0.5*0.5", degree=0)
#y_Omega = Expression("(x[0]-0.5)*(x[0]-0.5) + (x[1]-0.5)*(x[1]-0.5)", degree=2)
y_Omega =  Constant(2.0)


# variational statement for the coupled problem
a = inner(grad(y_),grad(v_))*dx + 1/lam*inner(beta*beta*p_,v_)*dx + \
    inner(grad(p_),grad(q_))*dx -       inner(          y_,q_)*dx
l = inner(-y_Omega,q_)*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(Y2.sub(0),y0,DirichletBoundary())
bc0 = DirichletBC(Y2.sub(1),Constant(0.0),DirichletBoundary())

# assemble and solve the coupled problem
yp = Function(Y2)
solve(a==l,yp,[bc,bc0])
y,p=yp.split()

# finally evaluate the optimal control
u  = Function(U)
project(-1/lam*beta*p,U,function=u)


# save the state, the adjoint state, and the optimal control
File("u.pvd") << u
File("y.pvd") << y
File("p.pvd") << p

# assemble and print the objective functional
J = 0.5*inner(y-y_Omega,y-y_Omega)*dx + lam/2*inner(u,u)*dx
print("J = %e" % assemble(J))

# check the norm of the "gradient" beta*p + lam*u
print("||beta*p + lam*u|| = %e" % sqrt(assemble(inner(beta*p+lam*u,beta*p+lam*u)*dx)) )