from firedrake import *


m=20
Ly = 0.85
mesh = IntervalMesh(m, 0 , Ly)

# Time definitions
t=0.0
end=10.0



# Define Function space on our mesh.
# Initially we will use a continuous linear Lagrange basis

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


# Define timestep value
Dt = 0.01
dt = Constant(Dt)


# Define Crank Nicholson parameter
theta = 1.0

# Define Groundwater constants
m = 0.3
sigma = 0.8
Lc = 0.05
k = 1e-8
w = 0.1
R = 0.000125
nu = 1e-6
g = 9.81
alpha = k  / ( nu * m * sigma )

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

# 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

h = Function(V)
phi = TestFunction(V)

# Provide intial guess to non linear solve
h.assign(h_prev)

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

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, Expression('0.07'), 1)

# Space for ODE evaluation for hc

h_problem = NonlinearVariationalProblem( F , h )
h_solver = NonlinearVariationalSolver(h_problem, bcs = [bc1])

while (t < end):
	# First we increase time
	t += Dt

	# Print to console the current time
	print "Time is %g" % (t)

	# Use the solver and then update values for next timestep
	h_solver.solve()
	h_prev.assign(h)

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