# Benchmark problem 3 for Navier-Stokes 
# =======================

from firedrake import *

N = 50

M = PeriodicSquareMesh(N, N, 2)

V = VectorFunctionSpace(M, "CG", 2)
V_out = VectorFunctionSpace(M, "CG", 1)
W = FunctionSpace(M, "CG", 1)
Z = V * W
Z_out = V_out * W

up_out = Function(Z_out)
up = Function(Z)
u, p = split(up)
v, q = TestFunctions(Z)
u_ = Function(Z.sub(0))

Dt = 0.2*(1/N) # 0.2*mesh_size/U_max

x,y = SpatialCoordinate(M)

nu = Constant(1/10) #kinematic viscosity

ic = project(as_vector([cos(pi*x)*sin(pi*y),cos(pi*y)*sin(pi*x)]),Z.sub(0)) # initial condition for velocity 

init_p = Function(W).interpolate(-0.25*(cos(2*pi*x) + cos(2*pi*y))) # initial pressure

u_.assign(ic)

U = 0.5*(u + u_)

F =( inner(u,v) * dx - inner(u_,v) *dx +
    Dt *nu * inner(nabla_grad(U), nabla_grad(v)) * dx +
    Dt *inner(dot(U, nabla_grad(U)), v) * dx -
    Dt *p * div(v) * dx +
    Dt*div(U) * q * dx )


nullspace = MixedVectorSpaceBasis(Z, [Z.sub(0), VectorSpaceBasis(constant=True)])

outfile = File("test_run_P3.pvd")


t = 0.0
freq = 1
end = 1
while (t<=end):
      
      solve(F == 0, up, nullspace=nullspace, solver_parameters={"ksp_type": "gmres", "mat_type": "aij","pc_type": "lu","pc_factor_mat_solver_type": "mumps"})    
      
      u, p = up.split()

      if freq%5==0: #printing results every 5 time steps
         print("t=", round(t,2))
         u_exact = project(as_vector([exp(-2*t*nu*pi*pi)*cos(pi*x)*sin(pi*y),exp(-2*t*nu*pi*pi)*cos(pi*y)*sin(pi*x)]),V)
         L_2_error = errornorm(u_exact,u, mesh = M)
         print("error_L2_norm=",round(L_2_error,3))
         p.rename("Pressure")
         u_out = project(up, Z_out).split()[0]
         u_out.rename("Velocity")
         outfile.write(u_out, p)

      u_.assign(u)
      t += Dt
      freq +=1