# Code in this cell makes plots appear an appropriate size and resolution in the browser window
%matplotlib widget
%config InlineBackend.figure_format = 'svg'

import matplotlib.pyplot as plt

plt.rcParams['figure.figsize'] = (11, 6)

from firedrake import *

from firedrake.adjoint import *

continue_annotation()

mesh = Mesh("stokes-control.msh")

# NBVAL_IGNORE_OUTPUT
from firedrake.pyplot import quiver, streamplot, tricontourf, triplot

fig, axes = plt.subplots()
triplot(mesh, axes=axes)
axes.axis("off")
axes.set_aspect("equal")
axes.legend(loc="upper right");

V = VectorFunctionSpace(mesh, "Lagrange", 2)
Q = FunctionSpace(mesh, "Lagrange", 1)
W = V*Q

v, q = TestFunctions(W)
u, p = TrialFunctions(W)

nu = Constant(1)     # Viscosity coefficient

x, y = SpatialCoordinate(mesh)
u_inflow = as_vector([y*(10-y)/25.0, 0])

noslip = DirichletBC(W.sub(0), (0, 0), (3, 5))
inflow = DirichletBC(W.sub(0), interpolate(u_inflow, V), 1)
static_bcs = [inflow, noslip]

g = Function(V, name="Control")
controlled_bcs = [DirichletBC(W.sub(0), g, 4)]
bcs = static_bcs + controlled_bcs

a = nu*inner(grad(u), grad(v))*dx - inner(p, div(v))*dx - inner(q, div(u))*dx
L = Constant(0)*q*dx

w = Function(W)
solve(a == L, w, bcs=bcs, solver_parameters={"mat_type": "aij",
                                             "ksp_type": "preonly",
                                             "pc_type": "lu",
                                             "pc_factor_shift_type": "inblocks"})

# NBVAL_IGNORE_OUTPUT
u_init, p_init = w.subfunctions

fig, axes = plt.subplots(nrows=2, sharex=True, sharey=True)
streamlines = streamplot(u_init, resolution=1/3, seed=0, axes=axes[0])
fig.colorbar(streamlines, ax=axes[0], fraction=0.046)
axes[0].set_aspect("equal")
axes[0].set_title("Velocity")

contours = tricontourf(p_init, 30, axes=axes[1])
fig.colorbar(contours, ax=axes[1], fraction=0.046)
axes[1].set_aspect("equal")
axes[1].set_title("Pressure");

u, p = split(w)
alpha = Constant(10)

J = assemble(1./2*inner(grad(u), grad(u))*dx + alpha/2*inner(g, g)*ds(4))
m = Control(g)
Jhat = ReducedFunctional(J, m)

get_working_tape().progress_bar = ProgressBar

g_opt = minimize(Jhat)

# NBVAL_IGNORE_OUTPUT
fig, axes = plt.subplots()
arrows = quiver(g_opt, axes=axes, scale=3)
axes.set_aspect("equal")
axes.set_title("Optimised boundary value");

print("Jhat(g) = %.8g\nJhat(g_opt) = %.8g" % (Jhat(g), Jhat(g_opt)))

g.assign(g_opt)
w_opt = Function(W)
solve(a == L, w_opt, bcs=bcs, solver_parameters={"mat_type": "aij",
                                                 "ksp_type": "preonly",
                                                 "pc_type": "lu",
                                                 "pc_factor_shift_type": "inblocks"},
      annotate=False)

# NBVAL_IGNORE_OUTPUT
u_opt, p_opt = w_opt.subfunctions

fig, axes = plt.subplots(nrows=2, sharex=True, sharey=True)
streamlines = streamplot(u_opt, resolution=1/3, seed=0, axes=axes[0])
fig.colorbar(streamlines, ax=axes[0], fraction=0.046)
axes[0].set_aspect("equal")
axes[0].set_title("Optimized velocity")

contours = tricontourf(p_opt, 30, axes=axes[1])
fig.colorbar(contours, ax=axes[1], fraction=0.046)
axes[1].set_aspect("equal")
axes[1].set_title("Optimized pressure");

tape = get_working_tape()
tape.clear_tape()