""" This solver is in support of the paper 'Analysis of an interface stabilised
finite element method: The advection-diffusion-reaction equation', 
by Garth N. Wells.

It has been used against DOLFIN version 0.9.4 and FFC version 0.7.0. DOLFIN and 
FFC are freely available at http://www.fenics.org. 
"""

__author__    = "Garth N. Wells (gnw20@cam.ac.uk)"
__date__      = "2009-10-28"
__copyright__ = "Copyright (C) 2009 Garth N. Wells"
__license__   = "GNU GPL Version 3.0"

from dolfin import *

# Optimize compilation of the form
parameters.optimize = True

# Define boundary of domain
class Boundary(SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary

def solve(mesh, order, model):

    # Advective velocity
    V_b  = VectorFunctionSpace(mesh, "CG", order+6)
    b = model.b(V_b) 

    # Source function
    V_exact = FunctionSpace(mesh, "CG", order+6)
    f = model.f(V_exact)

    # Model parameters
    mu    = model.mu
    kappa = model.kappa
    alpha = model.alpha(mesh, order)

    # Create function spaces k = order. V_cg is restricted to fecets
    V_cg = FunctionSpace(mesh, "CG", order, "facet")
    V_dg = FunctionSpace(mesh, "DG", order)
    mixed = V_cg + V_dg

    # Dirichlet bc list
    bcs = model.bcs(V_cg)

    # Prescribed boundary flux
    hbc = model.hbc(V_cg)

    # Test and trial functions
    vhat, vh = TestFunction(mixed)
    uhat, uh = TrialFunction(mixed)

    # Mesh-related functions
    n = V_dg.cell().n
    h = CellSize(mesh)

    # ( dot(v, n) + |dot(v, n)| )/2.0  = an on outflow,  0 on inflow 
    an = (dot(b, n) + abs(dot(b, n)))/2.0

    # ( -dot(v, n) + |dot(v, n)| )/2.0  = 0 on outflow,  -an on inflow 
    an_min = (-dot(b, n) + abs(dot(b, n)))/2.0

    # Standard flux
    sigma = -b*uh + kappa*grad(uh) 

    # Interface flux
    sigma_a = -dot(b, n)*uh + an_min*(uhat - uh)
    sigma_d = kappa*grad(uh) + (alpha/h)*kappa*(uhat -uh)*n
    sigma_hat = sigma_a + dot(sigma_d, n)

    # Bilinear form
    a_local =  mu*vh*uh*dx + dot(grad(vh), sigma)*dx \
             - vh('+')*sigma_hat('+')*dS - vh('-')*sigma_hat('-')*dS \
             + dot(grad(vh)('+'), n('+'))*kappa*(uhat('+') - uh('+'))*dS \
             + dot(grad(vh)('-'), n('-'))*kappa*(uhat('-') - uh('-'))*dS \
             - vh*sigma_hat*ds \
             + dot(grad(vh), n)*kappa*(uhat - uh)*ds

    a_global = vhat('+')*sigma_hat('+')*dS + vhat('-')*sigma_hat('-')*dS \
             + vhat*sigma_hat*ds \
             + vhat*an*uhat*ds

    a = a_local + a_global

    # Linear form
    L = vh*f*dx + vhat*an_min*hbc*ds

    # Create variational problem and solve
    problem = VariationalProblem(a, L, bcs)
    Uh = problem.solve()

    # Exact solution
    u = model.u(V_exact)

    # Compute (L2)^2 for (u - uh)
    M = (Uh[1]-u)*(Uh[1]-u)*dx
    return sqrt(assemble(M, mesh=mesh))


# Hyperbolic example
class hyperbolic():
    b     = lambda self, V: Constant(V.mesh(), (4.0/5.0, 3.0/5.0))
    tmp0 =   "(cos((pi*(1+x[0])*(1+x[1])*(1+x[1])/8))*pi*(1+x[1])*(1+x[1])/8)"
    tmp1 =   "(cos((pi*(1+x[0])*(1+x[1])*(1+x[1])/8))*2*pi*(1+x[0])*(1+x[1])/8)"
    source = "(4.0/5.0)*tmp0 + (3.0/5.0)*tmp1 + 1 + sin((pi*(1+x[0])*(1+x[1])*(1+x[1])/8))"
    source = source.replace('tmp0', tmp0)
    source = source.replace('tmp1', tmp1)
    f =   lambda self, V: Expression(self.source, V = V)
    mu    = 1.0
    kappa = 0.0
    bcs   = lambda self, V: []
    hbc   = lambda self, V: Constant(V.mesh(), 1.0)
    u     = lambda self, V: Expression("1 + sin(pi*(1+x[0])*(1+x[1])*(1+x[1])/8)", V = V)
    alpha = lambda self, mesh, order: Constant(mesh, 0.0)

# Elliptic example
class elliptic():
    b     = lambda self, V: Constant(V.mesh(), (0.0, 0.0))
    f     = lambda self, V: Expression("2.0*pi*pi*sin(1.0*pi*x[0])*sin(1.0*pi*x[1])", V = V)
    mu    = 0.0
    kappa = 1.0
    bcs   = lambda self, V: [DirichletBC(V, Constant(V.mesh(), 0.0), Boundary())]
    hbc   = lambda self, V: Constant(V.mesh(), 0.0)
    u     = lambda self, V: Expression("sin(1.0*pi*x[0])*sin(1.0*pi*x[1])", V = V)
    alpha = lambda self, mesh, order: Constant(mesh, 4.0*order*order)


# Advection-diffusion example source function
class advection_diffusion_source(Expression):
    def __init__(self, mu, kappa, V=None):
        self.mu    = mu
        self.kappa = kappa
    def eval(self, values, x):
        u   =  sin(1.0*pi*x[0])*sin(1.0*pi*x[1])
        vx  = -exp(x[0])*(x[1]*cos(x[1]) + sin(x[1]))
        vy  =  exp(x[0])*(x[1]*sin(x[1]))
        ux  =  1.0*pi*cos(1.0*pi*x[0])*sin(1.0*pi*x[1])
        uy  =  1.0*pi*sin(1.0*pi*x[0])*cos(1.0*pi*x[1])
        uxx = -2.0*pi*pi*sin(1.0*pi*x[0])*sin(1.0*pi*x[1])
        values[0] = self.mu*u + vx*ux + vy*uy - self.kappa*uxx

# Advection-diffusion example
class advection_diffusion:
    b     = lambda self, V: Expression(("-exp(x[0])*(x[1]*cos(x[1]) + sin(x[1]))", \
                                        "exp(x[0])*(x[1]*sin(x[1]))"), V = V)
    mu    = 0.0
    kappa = 10.0
    f     = lambda self, V: advection_diffusion_source(self.mu, self.kappa, V = V)
    bcs   = lambda self, V: [DirichletBC(V, Constant(V.mesh(), 0.0), Boundary()) ]
    hbc   = lambda self, V: Constant(V.mesh(), 0.0)
    u     = lambda self, V: Expression("sin(1.0*pi*x[0])*sin(1.0*pi*x[1])", V = V)
    alpha = lambda self, mesh, order: Constant(mesh, 4.0*order*order)


## Choose model here
#model  = hyperbolic()
#model = elliptic()
model = advection_diffusion()

# Cells in each direction
N = [2, 4, 6, 8, 10, 12, 14, 16, 18, 20]

# Cell sizes
h = []
for n in N:
    h.append(2.0/n)

# Polynomial orders
orders = [1, 2, 3, 4, 5]

# Compute error for each mesh
error = []
for order in orders:
    e = []
    for n in N:
        mesh = Rectangle(-1.0, -1.0, 1.0, 1.0, n, n)
        e.append(solve(mesh, order, model))
    error.append(e)

    print "Order: ",order 
    print h
    print e

print "Computed errors"
for i, e in enumerate(error):
    print "e",i+1,"=", e