"""
Solve the constrained 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
            ua <= u <= ub in Omega

using the conditional gradient method.

Recall that the directional derivative of the objective functional is given by

J'h = (beta p + lam u,h)

where p is the solution to the adjoint problem:

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

from dolfin import *

N = 64  # 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)
y_ = TrialFunction(Y)
v_ = TestFunction(Y)


lam = Constant(1.0E00)
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)

ua = Constant( 0.0E+00)
ub = Constant( 5.0E-02)


# vectors for storing the state and adjoint state and control
y = Function(Y)
p = Function(Y)
u = Function(U)

# weak form of the state problem
a    = inner(grad(y_),grad(v_))*dx
l    = inner(beta*u,v_)*dx
ladj = inner(y-y_Omega,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())
bc0 = DirichletBC(Y,Constant(0.0),DirichletBoundary())


# let us initialize u to be between ua and ub
project((ua+ub)/2,U,function=u)


J = 0.5*inner(y-y_Omega,y-y_Omega)*dx + lam/2*inner(u,u)*dx
s = Constant(0.0)

# let g be the Riesz reprezentation of the derivative:
g = beta*p + lam*u

i = 0

v = Function(U)
w = Function(Y)
lv = inner(beta*v,v_)*dx
# the conditional gradient loop; see p.93 in Troltsch
while True:
    # solve direct and adjoint problems
    # S1
    solve(a==l,   y,bc )
    # assemble and print the objective functional
    print("i = %3d, J = %e" % (i, assemble(J)))
    # S2
    solve(a==ladj,p,bc0)
    # save the state, the adjoint state, and the optimal control
    File("u-%03d.pvd" %i) << u
    File("y-%03d.pvd" %i) << y
    File("p-%03d.pvd" %i) << p
    # S3
    # f'(u)v = min_[v_\in Uadm] f'(u) v_
    # i.e. if g>0 we set v = ua, if g<0 then v=ub, else 0.5*(ua+ub)
    #project(conditional(g>0,ua,conditional(g<0,ub,0.5*(ua+ub))),U,function=v)
    project(conditional(g>0,ua,conditional(g<0,ub,u)),U,function=v)
    # S4: compute w = y(v)
    solve(a==lv,w,bc)
    # S4: compute step length
    g1 = assemble(inner(y-y_Omega,w-y)*dx+lam*inner(u,v-u)*dx)
    g2 = assemble(inner(w-y,w-y)*dx + lam*inner(v-u,v-u)*dx)
    s.assign(max(0.0,min(1.0,-g1/g2)))
    # check whether we need to stop
    if sqrt(assemble(inner(s*(v-u),s*(v-u))*dx)) < 1.0E-04:
        print('Stop: control update is too small!')
        break
    # S5: update the control
    v.vector().axpy(-1.0,u.vector()) # v = v-u
    u.vector().axpy(float(s),v.vector()) # u = u + s*v
    i += 1