from firedrake import *

# Script to solve incompressible acoustic waves in three dimensions with rotation.
# c^2_0 and \rho_0 will be considered to be scaled to be equal to 1.



# Create mesh
# Current mesh is a unit box that is periodic in the y-direction.
m = 16
n = 16.

Nx = m
Ny = m
Nz = m

element_height =  1. /n

m = PeriodicRectangleMesh(Nx, Ny, 1, 1, direction="y", quadrilateral=quadrilateral)
mesh = ExtrudedMesh(m, layers = Nz, layer_height = element_height )

# Declare timestep
# T/20 -  Period of wave /20
dt = 0.0493785246 


# Declare initial and end time

t = 0.
end = dt

# Declare order of the basis in the elements

order_basis = 1

# Declare flux indicator function
theta = Constant(0.5)


# Declare function spaces on the mesh
# and create mixed function space
# for coupled Poisson Bracket.

V = VectorFunctionSpace(mesh, "DG", order_basis)
R = FunctionSpace(mesh, "DG", order_basis)

W = V*R




# Initial conditions

# Function space
w0 = Function(W)

# Interpolate expressions
u0,P0 = w0.split()

# Period of waves considered in test is sigma

sigma = 0.15717672547

u0.interpolate(Expression(["-((2*pi*pow(sigma,3))/(1-pow(sigma,2)))*(1+1/pow(sigma,2))*sin(2*pi*x[0])*sin(2*pi*x[1])*cos(2*pi*x[2])", "(-2*pi*sigma*cos(2*pi*x[0])+sin(2*pi*x[0])/(2*pi))*cos(2*pi*x[1])*cos(2*pi*x[2])", "(sigma*cos(2*pi*x[0])+sin(2*pi*x[0]))*sin(2*pi*x[1])*sin(2*pi*x[2])"], sigma = sigma  ))

P0.interpolate(Expression("-(pow(sigma,2))*(cos(2*pi*x[0]) + sin(2*pi*x[0])/sigma)*cos(2*pi*x[1])*cos(2*pi*x[2]) ",sigma = sigma))




# Declare angular velocity
Omega = Function(V)
Omega.interpolate(Expression(["0.0","0.0", "1.0"]))


# Set up internal boundary normal on mesh 
# Needed for numerical fluxes in bilinear form
n = FacetNormal(mesh)

# Define discrete divergence operator
def div_u(u, p):
	return (dot(u, grad(p)))*dx + (jump(p)*dot((u('-')*(1-theta)+u('+')*theta), n('-')))*(dS_h+dS_v)


# Project the initial velocity to be divergence free
# U_divfree = U + Udagger
# implies DIV U = - DIV Udagger


dFdPproject = TestFunction(R)

DIVU = - div_u ( u0, dFdPproject )

u_perp = TrialFunction(V)

DIVU_perp =   div_u ( u_perp, dFdPproject )

u_perp = Function(V)

# Direct solve with LU does not work as DIV is not a square matrix
# We also can't subsequently use the ilu or jacobi preconditioner for the system.
# GMRES also requires a square matrix, so we use lsqr to solve our projection problem.

solve(DIVU_perp == DIVU, u_perp, solver_parameters={'ksp_rtol': 1e-14, 'pc_type' : 'none', 'ksp_type': 'lsqr'}
   )

u0 += u_perp

# Print statement to verify initial velocity is divergence free

u_div = assemble( div_u( u0 , dFdPproject ) )

u_div_norm = norm (u_div,  norm_type="L2")

print "l2 norm of discrete divergence on initial velocity is " % u_div_norm

# Assemble initial energy
E0 = assemble ( 0.5*(inner(u0,u0) )*dx)

#print E0

# Output or print that the discrete divergence of the velocity is zero.
#print div(u0)

# Define trial and test functions on the mixed space

(u, P) = TrialFunctions(W)
(dFdu_vec, dFdP) = TestFunctions(W)

# Create bilinear form for linear solver
# Bilinear problem is of the form 
# a(u,v) = L(v)
# Using our Poisson bracket and an implicit midpoint rule
# we see that
# a(u,v) = u^{n+1}*dFdu_vec + P^{n+1)*dFdP - 0.5*dt*PB(u^{n+1}, P^{n+1})
# L(v) = u^{n}*dFdu_vec + P^{n)*dFdP + 0.5*dt*PB(u^{n}, P^{n})

# We note that there are no boundary surface integrals ds, as we require
# the normal of the variational derivative and test function to vanish 
# at the boundary.

#Define varitional derivatives

(u0,P0) = split(w0)
dHdu0 = u0
dHdu = u

# Note that we project that DIV (U^n) = 0 
# So we only solve the equation DIV ( U ^n+1) = 0
# This will form a Poisson equation implicitly for the Pressure.

L0 = dot(u0, dFdu_vec) * dx
L1 = -dot(cross(2*Omega, dHdu0), dFdu_vec) * dx

L = L0 + 0.5 * dt * (L1  )


a0 = dot(u, dFdu_vec) * dx
a1 = -dot(cross(2*Omega, dHdu), dFdu_vec) * dx
a2 = -div_u(dFdu_vec, P)
a3 = div_u(dHdu, dFdP)

a = a0 - 0.5 * dt * ( a1 + a2  ) + a3

# Storage for visualisation
outfile = File('./Results/incompressible_acoustic_results.pvd')

u0,P0 = w0.split()

u0.rename("Velocity")
P0.rename("Pressure")


# Output initial conditions
outfile.write(u0,P0, time = t)


out = Function(W)
# File for energy output
E_file = open('./Results/energy.txt', 'w')


# Solve loop
E0 = None

while (t < end):
   # Update time
   t+= dt
 
   solve(a == L, out, nest=False,  solver_parameters={  "ksp_rtol": 1e-14, "pc_type": "ilu", })
   u, P = out.split()

   # Assign appropriate name in results file
   u.rename("Velocity")
   P.rename("Pressure")
 
   # Output results
   outfile.write(u, P, time =t)
 
   # Assign output as previous timestep for next time update
   u0.assign(u)
   P0.assign(P)

   # Assemble initial energy
   E = assemble ( (0.5*(inner(u,u) ))*dx)
   E0 = E0 or E
 
   E_file.write('%-10s %-10s\n' % (t,abs((E-E0)/E0)))
   # Print time and energy drift, drift should be around machine precision.
   print "At time %g, energy drift is %g" % (t, E-E0)



# Create analytic solutions for error analysis
exact_P= Function(R)
exact_P.interpolate(Expression("-(pow(sigma,2))*(cos(2*pi*x[0]) + sin(2*pi*x[0])/sigma)*cos(2*pi*x[1]- sigma*t)*cos(2*pi*x[2]) ",sigma = sigma, t=t))

exact_W= Function(R)
exact_W.interpolate(Expression("(sigma*cos(2*pi*x[0])+sin(2*pi*x[0]))*sin(2*pi*x[1] - sigma*t)*sin(2*pi*x[2])",sigma = sigma, t=t))

# Print error for Pressure
error_P = errornorm(P, exact_P,  norm_type='L2')
print "At time %g, l2 error in pressure is %g" % (t, error_P)

# Print error for vertical velocity
error_W = errornorm(u[2], exact_W,  norm_type='L2')
print "At time %g, l2 error in vertical velocity is %g" % (t, error_W)


# Close energy write
E_file.close()