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

from dolfin import *

parameters["ghost_mode"] = "shared_facet"


# Create mesh and define function space
N = 40
mesh = UnitSquareMesh(N,N)

# Test case
u_exact = Expression("cos(pi*x[0])*cos(2*pi*x[1])", degree=degree+2)
f_exact = Expression("cos(pi*x[0])*cos(2*pi*x[1])*5*pi*pi", degree=degree+2)


# Trial spaces
U_element = FiniteElement("DG", mesh.ufl_cell(), degree) # solution
S_element = VectorElement("DG", mesh.ufl_cell(), degree) # gradient
Uh_element= FiniteElement("CG", mesh.ufl_cell(), degree) #traces
Sh_element= FiniteElement("RT", mesh.ufl_cell(), degree) #fluxes
#Uh_element= FiniteElement("CG", mesh.ufl_cell(), degree+1, restriction="facet") #traces - This is not possible in the latest version of FENICS... use lowest order degree=1
#Sh_element= FiniteElement("RT", mesh.ufl_cell(), degree+1, restriction="facet") #fluxes
# Test spaces
V_element = FiniteElement("DG", mesh.ufl_cell(), degree+delta_p)
T_element = VectorElement("DG", mesh.ufl_cell(), degree+delta_p)
#
E = FunctionSpace(mesh,MixedElement([U_element,S_element,Uh_element,Sh_element,V_element,T_element]))

(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)+tau1,grad(v2)+tau2)*dx + inner(div(tau1),div(tau2))*dx + inner(v1,v2)*dx + inner(tau1,tau2)*dx
    # This is the "natural" inner prodict in the test space H1 x Hdiv, but the inner product above (related to the equation we are solving) works better
    #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):
    sighatn = dot(sighat,n)
    taun    = dot(tau, n)

    # eq1:  -div(sigma) = f
    eq1 = inner(sigma,grad(v))*dx - inner(sighatn('+'),v('+')-v('-'))*dS - inner(dot(sighat,n),v)*ds
    # eq2:  sigma - grad(u) = 0
    eq2 = inner(sigma,tau)*dx + inner(u,div(tau))*dx - inner(uhat('+'),taun('+')+taun('-'))*dS - inner(uhat,taun)*ds
    
    return eq1+eq2

# full mixed form
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)))

comm = mpi_comm_world()
mpiRank = MPI.rank(comm)

if mpiRank == 0:
    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