"""
Convergence test for mixed Poisson sphere example using L^2 x H^1
formulation, forcing the vector field into the tangent space via a
Lagrange multiplier.
"""

# Copyright (C) 2013 Colin J. Cotter, Andrew McRae
# Modified by Marie E. Rognes (meg@simula.no)

from dolfin import *
import os

parameters["form_compiler"]["cpp_optimize"] = True
parameters["form_compiler"]["optimize"] = True

# Create and initialize application parameters
application_parameters = Parameters("Discretization")
application_parameters.add("order", 1)
application_parameters.parse()
info(application_parameters, True)

order = application_parameters["order"]
errorfilename = "results/errors_DG%dxCG%d.py" % (order, order+1)
if not os.path.isdir("results"):
    os.mkdir("results")

# Define meshes
meshes = ["sphere_ico3", "sphere_ico4", "sphere_ico5", "sphere_ico6"]

def create_function_spaces(mesh, order):
    S = VectorFunctionSpace(mesh, "DG", order)
    L = FunctionSpace(mesh, "DG", order)
    V = FunctionSpace(mesh, "Lagrange", order+1)
    R = FunctionSpace(mesh, "R", 0)
    return (S, L, V, R)

def main():

    # Set-up some storage for the errors. NB: Storing one error at a time
    # in case things fail
    info_blue("Storing errors to %s" % errorfilename)
    errorfile = open(errorfilename, "w")
    errorfile.write("# Results from  from L^2 x H^1 mixed-poisson test case.\n")
    errorfile.write("L2_error_u = [0.0 for i in range(%d)];\n" % len(meshes))
    errorfile.write("H1_error_u = [0.0 for i in range(%d)];\n" % len(meshes))
    errorfile.write("h_max = [0.0 for i in range(%d)];\n" % len(meshes))

    # Iterate over meshes
    for (i, meshid) in enumerate(meshes):

        meshname = "../hdiv-l2/meshes/%s.xml.gz" % meshid
        info_blue("Computing on %s" % meshname)
        mesh = Mesh(meshname)

        # Define function spaces
        (S, L, V, r) = create_function_spaces(mesh, order)
        W = MixedFunctionSpace((S, L, V, r))

        # Force
        g = Expression("x[0]*x[1]*x[2]")

        # Normal direction
        up_expr = Expression(("x[0]","x[1]","x[2]"))
        up = interpolate(up_expr, S)

        (sigma, l, u, r) = TrialFunctions(W)
        (w, gamma, v, t) = TestFunctions(W)

        a = (dot(sigma,w) - dot(grad(u),w) - dot(w,up)*l
             + dot(sigma,grad(v)) + gamma*dot(sigma,up) + u*t + r*v)*dx
        L = -v*g*dx

        # Solve problem
        w = Function(W)
        solve(a == L, w)

        # Store solutions to .xml
        info_blue("Storing solutions")
        file = File("results/w_%s_CG%d.xml.gz" % (meshid, order))
        file << w

        # Exact solution
        u_exact = Expression("-x[0]*x[1]*x[2]/12")
        sigma, l, u, r = w.split()

        # Compute errors
        info_blue("Computing errors")
        L2_error_u = errornorm(u_exact, u, "L2", degree_rise=1)
        H1_error_u = errornorm(u_exact, u, "H10", degree_rise=1)

        # Store errors
        info_blue("Storing errors")
        errorfile.write("h_max[%d] = %g;\n" % (i, mesh.hmax()))
        errorfile.write("L2_error_u[%d] = %g;\n" % (i, L2_error_u))
        errorfile.write("H1_error_u[%d] = %g;\n" % (i, H1_error_u))
        errorfile.flush()

    errorfile.close()


if __name__ == "__main__":
    main()
