"""This demo program solves the coupled system

    - div grad y + \lambda^{-1} beta^2 p = 0
    - div grad p - y                     = -y_Omega

on the unit square with homogeneous Dirichlet boundary conditions

    y = 0
    p = 0

and then evaluates the control
    u = -\lambda^{-1} beta * p
"""

from dolfin import *

# Create mesh and define function space
mesh = UnitSquareMesh(50, 50)
H = FunctionSpace(mesh, "Lagrange", 1)
# Create a product of function spaces
H2 = MixedFunctionSpace([H,H])

# Define Dirichlet boundary
def boundary(x,on_boundary):
    # in other cases one may use something depending on x[0] and x[1]
    return on_boundary

# Define boundary conditions
# on the first sub-component of the solution (y)
bc0 = DirichletBC(H2.sub(0), Constant(0.0), boundary)
# and the second one (p)
bc1 = DirichletBC(H2.sub(1), Constant(0.0), boundary)


# Other problem parameters, e.g.
lam     = Constant(1.0)
beta    = Constant(1.0)
y_Omega = Expression("10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)")

# Define variational problem
y,p = TrialFunctions(H2)
z,q = TestFunctions(H2)
# first equation
a0 = inner(grad(y),grad(z))*dx + 1.0/lam*beta*beta*p*z*dx
# second equation
a1 = inner(grad(p),grad(q))*dx - y*q*dx
# full bilinear form
a = a0 + a1
# right hand side/linear functional
L = -y_Omega*q*dx

# Compute solution
yp_ = Function(H2)
solve(a == L, yp_, [bc0,bc1])

# Split the solution into the individual components
y_,p_=yp_.split()


# Compute the control
u_ = -1.0/lam * beta * p_
# In order to use it we should project it on some piecewise-polinomial space;
# we could e.g. use H here
u_ = project(u_,H)

# Save solution in VTK format: can be viewed with e.g. Paraview
#file = File("ex4_u.pvd")
#file << u_

# Plot the solution
plot(u_, interactive=True)