# 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()

# Define a simple mesh
n = 32
mesh = UnitSquareMesh(n, n)

# Define P2 function space and corresponding test function
V = VectorFunctionSpace(mesh, "Lagrange", 2)
v = TestFunction(V)

# Create Functions for the solution and time-lagged solution
u = Function(V, name="Velocity")
u_ = Function(V)

# Assign initial condition
x, y = SpatialCoordinate(mesh)
u_ = interpolate(as_vector([sin(pi*x), 0]), V)
u.assign(u_)

# Set diffusivity constant
R = FunctionSpace(mesh, "R", 0)
nu = Function(R)
nu.assign(0.0001)

# Define nonlinear form
dt = 1.0/n
F = (inner((u - u_)/dt, v) + inner(dot(u, nabla_grad(u)), v) + nu*inner(grad(u), grad(v)))*dx

# Timestepping details
end_time = 0.5
timesteps_per_export = 4
num_timesteps = int(end_time/dt)
num_exports = num_timesteps//timesteps_per_export

# Store forward solution at exports so we can plot again later
forward_solutions = [u.copy(deepcopy=True)]

# Time integrate
i = 0
t = 0.0
while (t < end_time):
    solve(F == 0, u)
    u_.assign(u)
    t += dt
    i += 1
    if i % timesteps_per_export == 0:
        forward_solutions.append(u.copy(deepcopy=True))

from firedrake.pyplot import tricontourf

fig, axs = plt.subplots(num_exports+1, sharex='col')

for i in range(len(forward_solutions)):
    tricontourf(forward_solutions[i], axes=axs[i])
    axs[i].annotate('t={:.2f}'.format(i*timesteps_per_export*dt), (0.05, 0.05), color='white');
axs[0].set_title('Forward solution, u');

J_form = inner(u, u)*ds(2)
J = assemble(J_form)
print("Objective value: {:.4f}".format(J))

stop_annotating();

tape = get_working_tape()
for i, block in enumerate(tape._blocks):
    print("Block {:2d}: {:}".format(i, type(block)))

g = compute_gradient(J, Control(nu))
print("Gradient of J w.r.t. diffusivity = {:.4f}".format(*g.dat.data_ro))

from firedrake.adjoint import get_solve_blocks

solve_blocks = get_solve_blocks()
assert len(solve_blocks) == num_timesteps

for block in solve_blocks:
    assert block.adj_sol is not None

# 'Initial condition' for both adjoint
dJdu = assemble(derivative(J_form, u))

# Plot adjoint solutions at matching timesteps to forward
fig, axs = plt.subplots(num_exports+1, 2, sharex='col', figsize=(10, 5))
for i in range(num_exports+1):
    t = i*timesteps_per_export*dt
    tricontourf(forward_solutions[i], axes=axs[i, 0])
    adjoint_solution = dJdu if i == num_exports else solve_blocks[timesteps_per_export*i].adj_sol
    # Get the Riesz representer
    adjoint_solution = dJdu.riesz_representation(riesz_map="H1")
    tricontourf(adjoint_solution, axes=axs[i, 1])
    axs[i, 0].annotate('t={:.2f}'.format(t), (0.05, 0.05), color='white');
    axs[i, 1].annotate('t={:.2f}'.format(t), (0.05, 0.05), color='white');

# Plot formatting
axs[0, 0].set_title('Forward solution');
axs[0, 1].set_title('Adjoint solution');