# Solve Poisson problem on unit square using ultraweak/discontinuous Petrov-Galerkin method

from dolfin import *

# Test case
u_exact = Expression("sin(pi*x[0])*sin(2*pi*x[1])")
f_exact = Expression("sin(pi*x[0])*sin(2*pi*x[1])*5*pi*pi")

# Create mesh and define function space
degree  = 1
delta_p = 2
N = 20
mesh = UnitSquareMesh(N,N)

# Trial spaces
U = FunctionSpace(mesh, "DG", degree)                         # solution
S = VectorFunctionSpace(mesh, "DG", degree)                   # gradient
Uh= FunctionSpace(mesh, "CG", degree+1, restriction="facet")  # traces
Sh= FunctionSpace(mesh, "RT", degree, restriction="facet")    # fluxes
# Test spaces
V = FunctionSpace(mesh, "DG", degree+delta_p)
T = VectorFunctionSpace(mesh, "DG", degree+delta_p)
#
E = MixedFunctionSpace([U,S,Uh,Sh,V,T])
#

(u,sig,uhat,sighat,v,tau)       = TrialFunctions(E)
(du,dsig,duhat,dsighat,dv,dtau) = TestFunctions(E)
n = FacetNormal(mesh)

# bilinear forms
# inner product in the test space
def iprod(v1,tau1,v2,tau2):
    return inner(grad(v1),grad(v2))*dx + inner(div(tau1),div(tau2))*dx + inner(v1,v2)*dx + inner(tau1,tau2)*dx
# linear differential operator
def b(u,sigma,uhat,sighat,v,tau):
    return inner(sigma,grad(v))*dx + inner(sigma,tau)*dx + inner(u,div(tau))*dx - inner(uhat('+'),dot(tau('+'),n('+'))+dot(tau('-'),n('-')))*dS - inner(uhat,dot(tau,n))*ds - inner(dot(sighat('+'),n('+')),jump(v))*dS - inner(dot(sighat,n),v)*ds 

a = iprod(v,tau,dv,dtau) + b(u,sig,uhat,sighat,dv,dtau) + b(du,dsig,duhat,dsighat,v,tau)

# linear functional
L = inner(f_exact,dv)*dx

# boundary conditions on uhat
bcs = []
bcs.append(DirichletBC(E.sub(2), u_exact, DomainBoundary()))

# Solve the system
uSol = Function(E)
solve(a==L, uSol, bcs)
(u,sigma,uhat,sighat,v,tau) = uSol.split()

# Estimate the error
L2_err = sqrt(assemble((u-u_exact)*(u-u_exact)*dx))
V_err  = sqrt(assemble(iprod(v,tau,v,tau)))
print 'L2 error = ', L2_err, ' V-error = ', V_err

# Save the solution
u_file = File("u.pvd",'compressed')
u_file << u
s_file = File("sigma.pvd",'compressed')
s_file << sigma