# Simple Helmholtz equation
# =========================
#
from firedrake import *
import sys

# choose basis and order
try:
    space = sys.argv[1]
except IndexError:
    space = "Lagrange"

if space == "Lagrange":
    try:
        order = int(sys.argv[2])
    except IndexError:
        order = 1
elif space == "Hermite":
    order = 3
else:
    order = 0

# create base mesh
base_mesh = UnitSquareMesh(10, 10)

# we will interpolate mesh parameterisation using an L2 projection
Vc = VectorFunctionSpace(base_mesh, space, order)
X = TrialFunction(Vc)
phi = TestFunction(Vc)

xc, yc = SpatialCoordinate(base_mesh)
fc = as_vector([xc, yc])

a_l2 = dot(phi, X) * dx
L_l2 = dot(fc, phi) * dx

X = Function(Vc)
solve(a_l2 == L_l2, X)

# create parametrised mesh
mesh = Mesh(X)
print(f'{space} ({order}) mesh with {mesh.num_cells()} elements created successfully.')


# solve problem on new mesh
V = FunctionSpace(mesh, 'Lagrange', 1)

u = TrialFunction(V)
v = TestFunction(V)

x, y = SpatialCoordinate(mesh)
f = (1+8*pi*pi)*cos(x*pi*2)*cos(y*pi*2)

a = (dot(grad(v), grad(u)) + v * u) * dx
L = f * v * dx
print('problem on mesh defined successfully')

u = Function(V)
solve(a == L, u)

print('problem solved successfully')

u_exact = cos(x*pi*2)*cos(y*pi*2)
print('error', sqrt(assemble(dot(u-u_exact, u-u_exact) * dx)))