# 
# 2D shallow water equation finite volume
# 
# step up (0): 1 linear advection equation with forward Euler
# step up (i): 3 linear advection equations with forward Euler
# step up (ii): Lax-Friedrich fluxes 
# step up (iii): fluxes 
# 
# 
from firedrake import *
m = 20
tmesh = Mesh("rectmesh.msh")
tend  = 4.211
timestep = tend/40

V = FunctionSpace(tmesh, "DG", 0)
W = VectorFunctionSpace(tmesh, "CG", 1)
V_out = FunctionSpace(tmesh, "CG", 1)

# Advection velocity
velocity = Expression(("1.0",  "0.0"))
u0 = project(velocity, W)

# Inflow BC
# inflow  = Expression("(x[0] < 0.05) ? 1.0 : 0") # why does this work
inflow  = Expression("(x[0] < 0.5) ? 1.0 : 0") # but this does not? because below it is interpolated in the entire domain and when the spatial mesh size/step is too small it projects to zero before the quadrature point?
# Inflow Variable
D0 = Function(V)
D0.interpolate(inflow)

# Initial Condition
ic = Expression("(x[0] < 0) ? 1.0 : 0")

# Set up element normal
n = FacetNormal(tmesh)
un = 0.5*(dot(u0, n) + abs(dot(u0, n)))

D_  = Function(V)
D  = TrialFunction(V)
phi = TestFunction(V)

# D.interpolate(ic)
D_.interpolate(ic)
# D_.assign(ic)

# 
# Setting up Bilinear form
# 
# 
thet = Constant(0.);
# 
a1 = ( (D/timestep)*phi - (1-thet)*D*dot(u0, grad(phi)) )*dx      # (1-thet) interior term at new time step 
a2 = (1-thet)*dot(jump(phi), un('+')*D('+') - un('-')*D('-'))*dS  # Internal flux at future time (1-thet)
a3 = (1-thet)*dot(phi, un*D)*(ds(112))                            # outflow at right wall, inflow at left wall
a  = a1 + a2 + a3
# 
L1 = (D_/timestep)*phi*dx+thet*( D_*dot(u0, grad(phi)) )*dx        # previous timestep 
# L2 = thet*dot(jump(phi), un('+')*D_('+') - un('-')*D_('-'))*dS   # Internal flux at old time weighted at thet
L2 = 0
L3 = -thet*dot(phi, un*D_)*(ds(112)) -D0*phi*dot(u0, n)*(ds(114))  # outflow at top wall & given inflow at left wall
L = L1+L2+L3 
F = a-L

out = Function(V)
DD = Function(V)
outfile = File("./Results/swe2d.pvd")
#  outfile.write(project(D, V_out, name="Velocity"))
# Write IC to file
# outfile << D_
# worked: 
D_.assign(out)
# worked: 
outfile.write(out) 
# outfile.wrote(project(out, W, name="D"))
t   = 0.0
while (t <= tend):
       #Solve problem
       solve(a == L, out) #  worked with: solve(a == L, out)
       # solve(F == 0, out)
       # Update previous timestep
       D_.assign(out)
	   #Update time
       t+=timestep
       #  outfile << out
       # worked with: 
       outfile.write(out) 
       # outfile.write(project(out, V_out, name="Velocity"))
       # outfile.write(project(out, W, name="D"))