from firedrake import *
mesh = BoxMesh(10, 10, 2, 2.0, 2.0, 0.1) 
V = FunctionSpace(mesh, 'Lagrange', 1)
p = Function(V, name="p")
phi = Function(V, name="phi")
u = TrialFunction(V)
v = TestFunction(V)
f = Function(V).interpolate(Expression("x[0]*x[1]*x[2]*t"))
outfile = File("waveEqn1.pvd")
outfile << phi
bcval = Constant(0.0)
bc = DirichletBC(V, bcval, 1)
T = 10.
dt = 0.001
theta = 0.5
t = 0
n = 0
step = 0
lump_mass = True
while t <= T:
    n += 1
    step += 1
    bcval.assign(sin(2*pi*5*t))
    phi -= dt * theta * p[n] + dt * (1 - theta) * p[n-1]
    if lump_mass:
        p += assemble(dt * (theta * (inner(nabla_grad(v), nabla_grad(phi[n]))*dx) + f[n]*v*dx) + \
                      ((1-theta) * (inner(nabla_grad(v), nabla_grad(phi[n-1]))*dx) + f(n-1)*v*dx)) / assemble(v*dx)
        bc.apply(p)
    else:
        solve(u * v * dx == v * p * dx + dt * (theta * (inner(grad(v), grad(phi[n])) * dx + f[n]*V*dx)+ \
                                               (1-theta) * (inner(grad(v), grad(phi[n-1])) * dx) + f[n-1]*v*dx),\
              p, bcs=bc, solver_parameters={'ksp_type': 'cg','pc_type': 'sor','pc_sor_symmetric': True})
    phi -= dt * theta * p[n] + dt * (1 - theta) * p[n-1]
    t += dt
    if step % 10 == 0:
        outfile << phi