'''
Forced convection problem
'''

from fenics import *
import mshr

#
# initialize MPI for running in parallel, e.g.
# mpirun -np 2 python convection.py
#
comm    = mpi_comm_world()
mpiRank = MPI.rank(comm)
parameters["ghost_mode"] = "shared_facet"

# Hint: to turn off the "solving variational problem" stuff uncomment the following line
set_log_level(WARNING)

#
# domain and mesh
#
W,H = 1,2
domain_vertices = [Point(0,0), Point(W,0), Point(W,H), Point(0,H)]
domain          = mshr.Polygon(domain_vertices)
# mesh resolution: higher numbers give more elements
mesh_resolution = 50
mesh            = mshr.generate_mesh(domain,mesh_resolution)
#
ncells = mesh.size_global(2)
if mpiRank==0:
    print("No. of cells:  %7d" % (ncells))

# define state space
# Taylor-Hood element for discretizing Stokes
E1 = VectorElement('CG', mesh.ufl_cell(), 2)
E2 = FiniteElement('CG', mesh.ufl_cell(), 1)
# Lagrange element for discretizing advection-diffusion
E3 = FiniteElement('CG', mesh.ufl_cell(), 1)
Y  = FunctionSpace(mesh,MixedElement([E1,E2,E3]))
# discontinuous Lagrange element for discretizing controls
EU = FiniteElement('DG', mesh.ufl_cell(), 0)
U  = FunctionSpace(mesh,EU)

y = Function(Y)
y_= TrialFunction(Y)
z_= TestFunction(Y)

yU,yP,yT = split(y)
U_,P_,T_ = split(y_)
V_,Q_,S_ = split(z_)

# gravity vector
eg = Constant((0.,-1.))
# Prandtl number (ratio of momentum diffusivity to thermal diffusivity)
Pr = Constant(1.0E-01)
# Grashof (the ratio of the buoyancy to viscous force acting on a fluid)
Gr = Constant(5.0E+04)
# boundary conductivity
alpha = Constant(1.)
# outside boundary temperature
T0 = Constant(1.0)
# coefficient in front of control
beta = Constant(2.0E+01)
# initialize control
u    = Function(U)
project(Expression("0.5",element=EU),U,function=u)

# Boundary conditions:
#   no-slip velocity on the boundary
#   zero reference pressure in one of the corners
noslip = Constant((0.0,0.0))
bc = [DirichletBC(Y.sub(0), noslip, "on_boundary"), \
      DirichletBC(Y.sub(1), Constant(0.), "x[0]<DOLFIN_EPS && x[1]<DOLFIN_EPS", method="pointwise")]

# Residual of the system
# 1. momentum equation
eq1 = -Pr*inner(grad(yU),grad(V_))*dx + inner(yP,div(V_))*dx  - Pr*Pr*Gr*inner(eg*yT,V_)*dx
# 2. continuity (incompressibility) equation
eq2 = -inner(div(yU),Q_)*dx
# 3. advection-diffusion equation
eq3 = -inner(grad(yT),grad(S_))*dx + inner(yT*yU,grad(S_))*dx - inner(alpha*(yT-T0),S_)*ds + inner(beta*u,S_)*dx
# add everything together
F = eq1 + eq2 + eq3

# Navier-Stokes can be included by adding
#F += inner(grad(U)*U,V_)*dx
# to F

# Compute the Jacobian of the system
jac = derivative(F,y,y_)

# Solve the non-linear system
solve(F==0,y,bc,J=jac,\
      solver_parameters={"newton_solver": \
                         {"relative_tolerance": 1e-12}})

# Write down the solutions
yu_file = File(comm,"yu.pvd","compressed")
yp_file = File(comm,"yp.pvd","compressed")
yT_file = File(comm,"yT.pvd","compressed")
u_file  = File(comm,"u.pvd", "compressed")

# save solutions at the initial iteration
iter = 0
yu_file << (y.sub(0),float(iter))
yp_file << (y.sub(1),float(iter))
yT_file << (y.sub(2),float(iter))
u_file  << (u,       float(iter))

# define the objective function
gamma = Expression(("1.0*(x[1]>=0.75*H)","0"),H=H,degree=1)
lmbda = Constant(1.0E+01)
yU_gamma = dot(gamma,y.sub(0)) 
J = 0.5*inner(gamma[0]*y.sub(0)[0],abs(y.sub(0)[0]))*dx \
  + 0.5*inner(gamma[1]*y.sub(0)[1],abs(y.sub(0)[1]))*dx + inner(0.5*lmbda*u,u)*dx

J0 = assemble(J)
if mpiRank==0:
    print("J = %e" % (J0))

# save the old control
u_old = Function(U)
u_old.vector().set_local(u.vector().array())
u_old.vector().apply('insert')
# form a random direction
from numpy import random, arange
h = Function(U)
h.vector().set_local(0.5-random.random(h.vector().local_size()))
h.vector().apply('insert')
    
#------------------------------------------------------------------------------
# Compute the derivative using finite differences
#------------------------------------------------------------------------------

# 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())
    u.vector().apply('insert')
    # resolve the system
    solve(F==0,y,bc,J=jac,\
      solver_parameters={"newton_solver": \
                         {"relative_tolerance": 1e-12}})
    # re-evaluate the functional
    J1 = assemble(J)
    # compute the FD approximation
    dJdh = (J1-J0)/delta
    if mpiRank==0:
        print("delta = %e, dJdh = %e" % (delta,dJdh))