from firedrake import *

m  = 80
Ly = 0.85
dy = Ly/m
mesh = IntervalMesh(m, 0 ,Ly)

# Time definitions
t   = 0.0
end = 150.0
Ntm = 400
dtmeas = end/Ntm
tmeas = dtmeas

# Define Function space on our mesh.

V = FunctionSpace(mesh, "CG", 1)

# Define timestep value

CFL = 0.3
Dt = CFL*0.5*dy*dy  # Based on FD estimate; note that dt must be defined before flux, etc 
# Don't understand why. Does this mean that for variable time step one need to redefine F again and again?
# dt.assign(CFL*0.5*dy*dy)

dt = Constant(Dt) # Using dt.assign in the while loop should avoid having the rebuild the solver iirc

# Define Crank Nicholson parameter
theta = 0.0

# Define Groundwater constants
mpor  = 0.3
sigma = 0.8
Lc    = 0.05
kperm = 1e-8
w     = 0.1
R     = 0.000125
nu    = 1.0e-6
g     = 9.81
alpha = kperm/( nu * mpor * sigma )
gam   = Lc/( mpor*sigma )
fac2  = sqrt(g)/( mpor*sigma )
# 
# ncase = 0 Dirichlet bc, ncase = 1 overflow groundwater into canal section with weir equation:
nncase = 1

# Initial condition
h_prev = Function(V)
h_prev.interpolate(Expression("0.00"))

# Create storage for paraview
outfile = File("./Results/groundwater.pvd")

# Write IC to file for paraview
outfile.write(h_prev , t = t )

# Define trial and test functions on this function space
# h will be the equivalent to h^n+1 in our timestepping scheme

out = Function(V)
phi = TestFunction(V)

def flux ( h , phi , R ):
	return ( alpha * g * h * dot ( grad (h) , grad (phi) ) - (R * phi )/ ( mpor * sigma ) )

## NB: Linear solves use TrialFunctions, non-linear solves use Functions with initial guesses.

if nncase == 0:
   # Provide intial guess to non linear solve

   h = Function(V)
   h.assign(h_prev)
   F = ( (h-h_prev)*phi/dt  + theta * flux ( h , phi , R ) + (1-theta)* flux ( h_prev, phi, R) ) *dx
   # Boundary conditions: Condition at Ly satisfied weakly
   bc1 = DirichletBC(V, 0.07, 1)
   h_problem = NonlinearVariationalProblem( F , h , bcs = bc1)

elif nncase == 1:

   if theta == 0.0:

     h = TrialFunction(V) # has to be set for linear solver
     aa = (h*phi/dt)*dx+(gam*phi*h/dt)*ds(1) # unknowns, only calculated once, not at each time step
     L2 = ( h_prev*phi/dt - flux ( h_prev, phi, R) ) *dx 
     L = L2+( gam*phi*h_prev/dt-phi*fac2*max_value(2.0*h_prev/3.0,0.0)*sqrt(max_value(2.0*h_prev/3.0,0.0)) )*ds(1)
     h_problem = LinearVariationalProblem(aa,L,out)
     h_solver = LinearVariationalSolver( h_problem)
   elif theta > 0.0:
     h = Function(V)
     h.assign(h_prev)

     F = ( (h-h_prev)*phi/dt  + theta * flux ( h , phi , R ) + (1-theta)* flux ( h_prev, phi, R) ) *dx
     # Add boundary contributions at y = 0: HERE HERE; does ds1 set it to the contribution at y=0????
     # A  F2 = ( gam*phi*(h-h_prev)/dt +theta*phi*fac2*max_value(2.0*h/3.0,0.0) )*ds(1)
     F2 = ( gam*phi*(h-h_prev)/dt+theta*phi*fac2*max_value(2.0*h/3.0,0.0)*sqrt(max_value(2.0*h/3.0,0.0))+(1-theta)*phi*fac2*max_value(2.0*h_prev/3.0,0.0)*sqrt(max_value(2.0*h_prev/3.0,0.0)) )*ds(1)
     h_problem = NonlinearVariationalProblem( F+F2 , h )
     h_solver = NonlinearVariationalSolver(h_problem)


while (t < end):
    # First we increase time
	t += Dt
    # Print to console the current time
    # Use the solver and then update valu`es for next timestep
	if theta == 0.0:
           h_solver.solve()
           h_prev.assign(out) # has to be renamed via out
        elif theta > 0.0:
           h_solver.solve()
           h_prev.assign(h)
        # Write output to file for paraview visualisation
        if t>tmeas:
            print "Time is %g" % (t)
            tmeas = tmeas+dtmeas
        
# End while loop
outfile.write(h_prev , t = t )



# implicit time step, aa only solved for once by using .assign command for h_prev