# 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 numpy import linspace

n = 100
mesh = PeriodicIntervalMesh(n, length=2)

x = SpatialCoordinate(mesh)[0]

u_init = sin(2*pi*x)

nu = Constant(1e-2)

V = FunctionSpace(mesh, "Lagrange", 2)

u_n1 = Function(V, name="u^{n+1}")
u_n = Function(V, name="u^{n}")
v = TestFunction(V)

u_n.interpolate(u_init)
dt = 1.0 / n

F = (((u_n1 - u_n)/dt) * v +
     u_n1 * u_n1.dx(0) * v + 
     nu*u_n1.dx(0)*v.dx(0))*dx

# If passed an existing Function object, the Function 
# constructor makes a copy.
results = [Function(u_n)]

t_end=0.5
for t in ProgressBar("Time step").iter(linspace(0.0, t_end, int(t_end/dt))):
    solve(F == 0, u_n1)
    u_n.assign(u_n1)
    results.append(Function(u_n))

help(plot)

# NBVAL_IGNORE_OUTPUT
from firedrake.pyplot import plot

fig, axes = plt.subplots()
axes.set_ylim((-1., 1.))
plot(results[0], axes=axes)

# NBVAL_IGNORE_OUTPUT
from matplotlib.animation import FuncAnimation

def animate(u):
    axes.clear()
    plot(u, axes=axes)
    axes.set_ylim((-1., 1.))

interval = 4e3 * float(dt)
animation = FuncAnimation(fig, animate, frames=results, interval=interval)

from IPython.display import HTML
HTML(animation.to_jshtml())

problem = NonlinearVariationalProblem(F, u_n1)

solver = NonlinearVariationalSolver(problem)

t = 0
t_end = 0.5
while t <= t_end:
    solver.solve()
    u_n.assign(u_n1)
    t += dt