from firedrake import *
from firedrake.petsc import PETSc
parameters['coffee']['O2'] = False

L = 1.0
H = 1.0
nlayers = 20
dz = H/nlayers

m = PeriodicIntervalMesh(nlayers, length=1.0)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=dz)

#1D spaces in horizontal
IUh0 = FiniteElement("CG", "interval", 1)
IUh1 = FiniteElement("DG", "interval", 0)

#1D spaces in vertical
IUv0 = FiniteElement("CG", "interval", 1)
IUv1 = FiniteElement("DG", "interval", 0)

#DG space
V2_elt = OuterProductElement(IUh1, IUv1)
V2 = FunctionSpace(mesh, V2_elt)

#Velocity space
V1_v = OuterProductElement(IUh1, IUv0)
V1_h = OuterProductElement(IUh0, IUv1)
V1_elt = HDiv(V1_v) + HDiv(V1_h)
V1 = FunctionSpace(mesh,V1_elt)

#Facet space and broken space for hybridisation
V1f = FunctionSpace(mesh,FacetElement(V1_elt))
V1b = FunctionSpace(mesh,BrokenElement(V1_elt))

V1t = FunctionSpace(mesh, V1_v)

#(u,w),v,p
W = MixedFunctionSpace((V1,V2,V1t,V2))

#rhs variable
F = project(Expression(("exp(x[1])*sin(x[0])+sinh(x[1])")),V2)

#Get F
u, v, b, p = TrialFunctions(W)
w, dv, gamma, q = TestFunctions(W)

#Boundary conditions
bcs = [DirichletBC(W[0], Expression(("0.", "0.")), "bottom"),
       DirichletBC(W[0], Expression(("0.", "0.")), "top")]
#bcs=[]

f = Constant(1.0)

a = (
    inner(w,u) - w[0]*f*v - div(w)*p + div(u)*q
    + dv*v + f*dv*u[0] + gamma*b + gamma*u[1]
    )*dx

L = q*F*dx

aP = (
    (inner(w,u) + div(w)*div(u))
    + p*q + dv*v + gamma*b
    )*dx

nullspace = MixedVectorSpaceBasis(W, [W.sub(0),
                                      W.sub(1),
                                      W.sub(2),
                                      VectorSpaceBasis(constant=True)
                                      ])


solver_parameters = {'ksp_type':'gmres',
                     'pc_type': 'fieldsplit',
                     'pc_fieldsplit_type': 'additive',
                     'ksp_converged_reason': True,
                     'ksp_monitor':True,
                     'fieldsplit_0_ksp_type': 'preonly',
                     'fieldsplit_1_ksp_type': 'preonly',
                     'fieldsplit_2_ksp_type': 'preonly',
                     'fieldsplit_3_ksp_type': 'preonly',
                     'fieldsplit_0_pc_type': 'lu',
                     'fieldsplit_1_pc_type': 'lu',
                     'fieldsplit_2_pc_type': 'lu',
                     'fieldsplit_3_pc_type': 'lu'}

w1 = Function(W)
uproblem = LinearVariationalProblem(a,L,w1,bcs=bcs,
                                    aP = aP)
usolver = LinearVariationalSolver(uproblem,nullspace=nullspace,
                                  solver_parameters=solver_parameters)
usolver.solve()
u,v, b,p = w1.split()
File('u1.pvd') << u