"""
0. Solve Laplace eqn using FENICs
1. Evaluate a given objective function
"""

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(0.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

print("J = %e" % assemble(J))

# save the solution as VTK file
file_y = File("y.pvd")
file_y << y