#
# Solve Laplace problem on unit square using ultraweak/discontinuous Petrov-Galerkin method
#
# Instead of solving the full mixed problem use a CG (conjugate gradient) iteration in the trial space
#

from dolfin import *
parameters["ghost_mode"] = "shared_facet"


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

# Test case
u_exact = Expression("sin(2*pi*x[0])*sin(3*pi*x[1])", degree=degree+2)
f_exact = Expression("sin(2*pi*x[0])*sin(3*pi*x[1])*13*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)
#
# Trial and test spaces
U=FunctionSpace(mesh,MixedElement([U_element,S_element,Uh_element,Sh_element]))
V=FunctionSpace(mesh,MixedElement([V_element,T_element]))
# Boundary conditions
bc  = DirichletBC(U.sub(2),u_exact,     "on_boundary")
bc0 = DirichletBC(U.sub(2),Constant(0), "on_boundary")


# inner products
# Trial space
def inner_U(u1,sigma1,uhat1,sighat1,u2,sigma2,uhat2,sighat2):
    sighat1n = dot(sighat1,n)
    sighat2n = dot(sighat2,n)
    return (inner(u1,u2)+inner(sigma1,sigma2))*dx \
         + (inner(uhat1('+'),uhat2('+')))*dS + inner(sighat1n('+'),sighat2n('+'))*dS \
         + (inner(uhat1,uhat2))*ds + inner(sighat1,sighat2)*ds \

# Test space
def inner_V(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


# bilinear form defining the 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

# and the right hand side
def ell(v,tau):
    return inner(f_exact,v)*dx


(u,sigma,uhat,sighat)       = TrialFunctions(U)
(du,dsigma,duhat,dsighat)   = TestFunctions(U)
(v,tau)                     = TrialFunctions(V)
(dv,dtau)                   = TestFunctions(V)
n                           = FacetNormal(mesh)

# vector for keeping the solution
X = Function(U)
Xu,Xsigma,Xuhat,Xsighat=X.split()

# prepare the solvers - we are going to reuse them
# Riesz mapping solver for U space
AU,BU     = assemble_system(inner_U(u,sigma,uhat,sighat,du,dsigma,duhat,dsighat),inner(Constant(0),du)*dx,bcs=bc)
AU_solver = LUSolver(AU)
AU_solver.parameters['reuse_factorization'] = True
# Riesz mapping solver for V space
AV,BV     = assemble_system(inner_V(v,tau,dv,dtau),ell(dv,dtau)-b(Xu,Xsigma,Xuhat,Xsighat,dv,dtau))
AV_solver = LUSolver(AV)
AV_solver.parameters['reuse_factorization'] = True

# start by initializing - uhat needs to satisfy the boundary conditions - this does not work particularly well, so
# we focus on the homogeneous problem instead.
#solve(inner_U(u,sigma,uhat,sighat,du,dsigma,duhat,dsighat)==inner(Constant(0),du)*dx,X,bc)
AU_solver.solve(X.vector(),BU)

# this is going to be an auxiliary function in V space, used to compute matrix-vector product
r = Function(V)
rv,rtau=r.split()

# residual in U
R = Function(U)
Ru,Rsigma,Ruhat,Rsighat=R.split()

# "search direction" oof CG method
P = Function(U)
Pu,Psigma,Puhat,Psighat=P.split()

# AP=A*P in U
AP = Function(U)
APu,APsigma,APuhat,APsighat=AP.split()

# see for example https://en.wikipedia.org/wiki/Conjugate_gradient_method#Example_code_in_MATLAB_.2F_GNU_Octave
# for details on CG iteration

# compute the initial residual
# 1. find the residual r in V
#solve(inner_V(v,tau,dv,dtau)==ell(dv,dtau)-b(Xu,Xsigma,Xuhat,Xsighat,dv,dtau),r)
AV_solver.solve(r.vector(),BV)
# 2. find the corresponding residual R in U
#solve(inner_U(u,sigma,uhat,sighat,du,dsigma,duhat,dsighat)==b(du,dsigma,duhat,dsighat,rv,rtau),R,bc0)
assemble(b(du,dsigma,duhat,dsighat,rv,rtau),tensor=BU)
bc0.apply(BU)
AU_solver.solve(R.vector(),BU)

# P = R
P.vector().axpy(1.0,R.vector())
# Rsold = ||R||_U
Rsold = assemble(inner_U(Ru,Rsigma,Ruhat,Rsighat,Ru,Rsigma,Ruhat,Rsighat))
Rs0   = Rsold

# main CG loop
k = 0
while k<100000:
    # AP = A*P
    #solve(inner_V(v,tau,dv,dtau)==b(Pu,Psigma,Puhat,Psighat,dv,dtau),r)
    assemble(b(Pu,Psigma,Puhat,Psighat,dv,dtau),tensor=BV)
    AV_solver.solve(r.vector(),BV)
    #solve(inner_U(u,sigma,uhat,sighat,du,dsigma,duhat,dsighat)==b(du,dsigma,duhat,dsighat,rv,rtau),AP,bc0)
    assemble(b(du,dsigma,duhat,dsighat,rv,rtau),tensor=BU)
    bc0.apply(BU)
    AU_solver.solve(AP.vector(),BU)
    # alpha = Rsold / (P' * AP)
    alpha = Rsold / assemble(inner_U(Pu,Psigma,Puhat,Psighat,APu,APsigma,APuhat,APsighat))
    # X = X + alpha * P
    X.vector().axpy( alpha,P.vector())
    # R = R - alpha * AP
    R.vector().axpy(-alpha,AP.vector())
    # Rsnew = ||R||_U
    Rsnew = assemble(inner_U(Ru,Rsigma,Ruhat,Rsighat,Ru,Rsigma,Ruhat,Rsighat))
    print('Iter: %6d,  rel.resid.: %e' % (k,sqrt(Rsnew/Rs0)))
    if Rsnew/Rs0 < 1.0E-12:
        break
    # P = R + (Rsnew / Rsold) * P
    beta = (Rsnew / Rsold)
    P.vector()[:] *= beta
    P.vector().axpy(1.0, R.vector())
    #
    Rsold = Rsnew
    k = k+1
# end while

# Estimate the error
L2_err = sqrt(assemble((Xu-u_exact)*(Xu-u_exact)*dx))
print 'L2 error = ', L2_err


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