"""
Convergence test for mixed Poisson sphere example.
"""

# Copyright (C) 2013 Marie E. Rognes, Colin J. Cotter, Andrew McRae

from dolfin import *
from ufl.log import info_blue
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("family", "RT")
application_parameters.add("order", 1)
application_parameters.add("exact_case", 2)
application_parameters.parse()
info(application_parameters, True)
family = application_parameters["family"]
order = application_parameters["order"]
exact_case = application_parameters["exact_case"]
errorfilename = "results/errors_%s_%d.py" % (family, order)
if not os.path.isdir("results"):
    os.mkdir("results")

# -----------------------------------------------------------------------------
info_green("Running with %s_%d and exact_case %d" % (family, order, exact_case))
# Define global normal
global_normal = Expression(("x[0]", "x[1]", "x[2]"))

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

# 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 H(div) x L^2 mixed-poisson test case.\n")
errorfile.write("L2_error_u = [0.0 for i in range(%d)];\n" % len(meshes))
errorfile.write("L2_error_sigma = [0.0 for i in range(%d)];\n" % len(meshes))
errorfile.write("Hdiv_error_sigma = [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 = "meshes/%s.xml.gz" % meshid
    info_blue("Computing on %s" % meshname)
    mesh = Mesh(meshname)
    mesh.init_cell_orientations(global_normal)

    # Space for exact solution
    Ve = VectorFunctionSpace(mesh, "Lagrange", order + 1)

    # Define approximation spaces and basis functions
    V = FunctionSpace(mesh, family, order)
    Q = FunctionSpace(mesh, "DG", order-1)
    R = FunctionSpace(mesh, "R", 0)
    W = MixedFunctionSpace((V, Q, R))
    (sigma, u, r) = TrialFunctions(W)
    (tau, v, t) = TestFunctions(W)

    # Choose exact solution based on user preference
    if (exact_case == 1):
        g = Expression("x[2]")
        u_exact = Expression("-0.5*x[2]")
    elif (exact_case == 2):
        g = Expression("x[0]*x[1]*x[2]")
        u_exact = Expression("-x[0]*x[1]*x[2]/12")
        sigma_exact = Expression(("-(x[1]*x[2]-3*x[0]*x[0]*x[1]*x[2]/(x[0]*x[0]+x[1]*x[1]+x[2]*x[2]))/12",
                                  "-(x[0]*x[2]-3*x[0]*x[1]*x[1]*x[2]/(x[0]*x[0]+x[1]*x[1]+x[2]*x[2]))/12",
                                  "-(x[0]*x[1]-3*x[0]*x[1]*x[2]*x[2]/(x[0]*x[0]+x[1]*x[1]+x[2]*x[2]))/12"))
    elif (exact_case == 3):
        g = Expression("3*x[2]*x[2] - 1")
        u_exact = Expression("-(3*x[2]*x[2] - 1)/6")
    else:
        raise Exception("Unrecognized exact case (%d)" % exact_case)

    # Define forms
    a = (inner(sigma, tau) + div(sigma)*v + div(tau)*u + r*v + t*u)*dx
    L = g*v*dx

    # Solve problem
    w = Function(W)
    solve(a == L, w)
    (sigma, u, r) = w.split()

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

    # Compute errors
    info_blue("Computing errors")
    L2_error_u = errornorm(u_exact, u, "L2", degree_rise=1)
    sigmadgmesh = VectorFunctionSpace(mesh,"DG",order)
    sigma_dg = project(sigma,sigmadgmesh)
    L2_error_sigma = assemble(inner(sigma_dg-sigma_exact,sigma_dg-sigma_exact)*dx)**0.5
    Hdiv_error_sigma = assemble(inner(sigma_dg-sigma_exact,sigma_dg-sigma_exact)*dx + (div(sigma_dg)-g)*(div(sigma_dg)-g)*dx)**0.5
    print "Sigma L2 = ", L2_error_sigma, "Sigma Hdiv = ",Hdiv_error_sigma
    # FIXME: Should also compute errors in L2 norm and H(div) seminorm
    # for sigma here too:

    # 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("L2_error_sigma[%d] = %g;\n" % (i, L2_error_sigma))
    errorfile.write("Hdiv_error_sigma[%d] = %g;\n" % (i, Hdiv_error_sigma))
    errorfile.flush()

errorfile.close()
info_green("Mission complete. See %s for error results." % errorfilename)
