#!/usr/bin/env python
# coding: utf-8

# In[15]:


# %load StokesSolver.py
#!/usr/bin/env python

#import netgen.gui
from ngsolve import *
from ngsolve.webgui import Draw
from netgen.geom2d import unit_square
from math import pi
import numpy as np


# ## Stokes Solver

# In[2]:


# In[2]:
# Turn off messages from ngsolve
ngsglobals.msg_level = 0

jump      = lambda v : v-v.Other()
avg       = lambda v : 0.5*(v+v.Other())
flux      = lambda v, n : 0.5*(Grad(v)+Grad(v).Other())*n
tang      = lambda v, n : v - (v*n)*n
tang_jump = lambda v, n: tang(v-v.Other(), n)

def StokesSolver(nu, u_D, f, p_ex=None,
                 *,                                           # From here, only keyword parameters allowed
                 space_u="H1", order_u=2, dgjumps_u = False,  # Space for u 
                 space_p="H1", order_p=1, dgjumps_p = False, lagrange_mult=False,  # Space for p             
                 init_maxh=1.0, init_refs=0, num_ref=0        # Mesh related parameters
                 ):

    # Check passed parameters for correctness
    valid_spaces_u = ["L2", "HDiv", "H1"]
    valid_spaces_p = ["L2", "H1"]
    if not space_u in valid_spaces_u:
        raise ValueError(f"{space_u} is not a valid velocity space, pick one from {valid_spaces_u}")
    if not space_p in valid_spaces_p:
        raise ValueError(f"{space_p} is not a valid pressure space, pick one from {valid_spaces_p}")

    # Generating mesh starting from only two cell to create completely structured mesh
    mesh = Mesh (unit_square.GenerateMesh(maxh=init_maxh))

    for ref in range(init_refs+num_ref):
          mesh.Refine()
    #Draw(mesh)
    
    # Pick velocity function space, for HDiv we must activate dgjumps
    # since we use classical DG approach and not HDG 
    V = (HDiv if space_u == "HDiv" else VectorL2 if space_u == "L2" else VectorH1)(mesh,order=order_u, dgjumps=dgjumps_u)
    Q = (L2 if space_p == "L2" else H1)(mesh, order=order_p, dgjumps=dgjumps_p)
    N = NumberSpace(mesh)
    if lagrange_mult:
        X = FESpace([V,Q,N], dgjumps=(dgjumps_u or dgjumps_p))  
        (u,p,lam), (v,q,mu) = X.TnT()
    else:
        X = FESpace([V,Q], dgjumps=(dgjumps_u or dgjumps_p))
        (u,p), (v,q) = X.TnT()
        lam, mu = None, None

    # Parameters in bilinear form
    gamma = Parameter(10)*order_u**2
    gamma_bnd = Parameter(10)*order_u**2
    beta_p = Parameter(0.1)*order_p**2
    n = specialcf.normal(2)
    h = specialcf.mesh_size

    A_h = BilinearForm(X)
    a = nu*InnerProduct(Grad(u),Grad(v))*dx
    # Nitsche terms for u
    a += nu*(-InnerProduct(Grad(u)*n, v) 
             -InnerProduct(u, Grad(v)*n)
             +gamma_bnd/h*InnerProduct(u, v))*ds(skeleton=True)

    a += (-div(u)*q-div(v)*p)*dx
    # Nitsche terms for p
    a += (u*n*q + v*n*p)*ds(skeleton=True)
    # Lagrange multiplicator to fix pressure
    if lagrange_mult:
        a += (lam*q+mu*p)*dx

    if space_u == "HDiv":
        a += nu*(-InnerProduct(flux(u,n),tang_jump(v,n)) 
                 -InnerProduct(tang_jump(u,n),flux(v,n)) 
                 +gamma/h*InnerProduct(tang_jump(u,n), tang_jump(v,n)))*dx(skeleton=True)

    if space_u == "L2":
        a += nu*(-InnerProduct(flux(u,n), jump(v)) 
                 -InnerProduct(jump(u),flux(v,n)) 
                 +gamma/h*InnerProduct(jump(u), jump(v)))*dx(skeleton=True)

    # Add discontinuous Galerkin term to b_h if velocity space is DG
    if space_u == "L2" and space_p == "L2":
        a += InnerProduct(jump(v), avg(p)*n)*dx(skeleton=True)
        a += InnerProduct(jump(u), avg(q)*n)*dx(skeleton=True)
    # Add pressure stabilization for discontinuous pressure if necessary
    if space_p == "L2" and (space_u == "H1" or (order_u == order_p)):
        a -= beta_p*jump(p)*jump(q)*dx(skeleton=True)

    A_h += a

    l = LinearForm(X)
    l += InnerProduct(f, v)*dx
    # Nitsche terms for rhs
    l += nu*(-InnerProduct(u_D, Grad(v)*n)
             +gamma_bnd/h*InnerProduct(u_D, v))*ds(skeleton=True)
    l += u_D*n*q*ds(skeleton=True)

    # Include correct mean value of p if given and Lagrange Multiplier is used
    if  lagrange_mult and p_ex:
        l += (mu*p_ex)*dx

    A_h.Assemble()
    l.Assemble()
    
    # Solve system
    ainv = A_h.mat.Inverse(freedofs=X.FreeDofs(), inverse="umfpack")
    upl = GridFunction(X)
    upl.vec.data = ainv*l.vec
    
    return upl.components[0], upl.components[1]


# ## Setting up a convergence test

# In[ ]:


nu = Parameter(1.0)

## Test example 1  
# u_ex = CoefficientFunction((sin(y), 0))
# p_ex = CoefficientFunction(sin(x))
# f = CoefficientFunction((sin(y) + cos(x),0)) 
# u_D = u_ex

## Test example 2
u_ex = CoefficientFunction((sin(pi*y), -cos(pi*x)))
u_D = u_ex
p_ex = CoefficientFunction(cos(pi*x)*cos(pi*y))
f = CoefficientFunction((-pi*sin(pi*x)*cos(pi*y) + nu * pi**2*sin(pi*y), -pi*sin(pi*y)*cos(pi*x) - nu * pi**2*cos(pi*x)))

# To collect error 
h1_errors_u = []
l2_errors_u = []
l2_errors_p = []
div_u_norms = []

# Number of refinements you want to do in the EOC study
max_num_ref = 2

solver_prms = {
    "space_u": "H1",
    "order_u": 2,
    "dgjumps_u": False,
    "space_p": "H1",
    "order_p": 1,   
    "dgjumps_p": False,
    "lagrange_mult" : True,
    "init_maxh": 0.3,
    "init_refs": 0,
    "num_ref" : 0,
}


# ## Running the convergence test

# In[ ]:


for num_ref in range(max_num_ref+1):
    u_h, p_h = StokesSolver(nu, u_D, f, p_ex, **solver_prms)
    mesh = u_h.space.mesh
    l2_errors_u.append((Integrate(InnerProduct(u_h-u_ex, u_h-u_ex), mesh))**0.5)
    l2_errors_p.append((Integrate(InnerProduct(p_h-p_ex, p_h-p_ex), mesh))**0.5)
    div_u_norms.append(Integrate(div(u_h)*div(u_h), mesh)**0.5)
    
    # Compute H1 error
    # N.B. H1 error cannot be computed directly.
    # We do this by interpolating u_ex into a higher order space
    order_ex = solver_prms["order_u"]+2
    Vh_ex = VectorH1(mesh, order=order_ex)
   
    # Interpolate u_ex into higher order space and compute derivative
    u_ex_h = GridFunction(Vh_ex)
    u_ex_h.Set(u_ex)
    grad_u_ex_h = u_ex_h.Deriv()  
    
    # Interpolate u_h into higher order space and compute derivative
    u_h_inter = GridFunction(Vh_ex)
    u_h_inter.Set(u_h)
    grad_u_h_inter = u_h_inter.Deriv()

    #h1_errors_u.append((Integrate(InnerProduct(grad_u_ex_h, grad_u_ex_h), mesh))**0.5)
    h1_errors_u.append((Integrate(InnerProduct(grad_u_h_inter-grad_u_ex_h, grad_u_h_inter - grad_u_ex_h), mesh, order=order_ex))**0.5)

    print("H1 error u = {:e}".format(h1_errors_u[-1]))
    print("L2 error u = {:e}".format(l2_errors_u[-1]))
    print("L2 error p = {:e}".format(l2_errors_p[-1]))
    print("Div(u) norm  = {:e}".format(div_u_norms[-1]))

    vtk = VTKOutput(ma=mesh,coefs=[u_ex, u_h, p_ex, p_h],
                    names=["u_ex","u_h","p_ex", "p_h"],
                    filename=f'StokesSolver_N_{solver_prms["num_ref"]}',
                    subdivision=2)
    vtk.Do()

    # Increase number of refinements
    solver_prms["num_ref"] += 1


# ## Postprocessing

# In[24]:


# Compute convergence rates, assuming uniform mesh refinement
# resulting in h reduction by a factor of 2
# Insert also 'inf' at beginning of eoc such that errs and eoc have same length 

compute_eoc = lambda errors : np.insert( (np.log(errors[:-1]) - np.log(errors[1:]))/np.log(2), 0, np.Inf)

h1_eocs_u = compute_eoc(h1_errors_u)
l2_eocs_u = compute_eoc(l2_errors_u)
l2_eocs_p = compute_eoc(l2_errors_p)
div_u_eocs = compute_eoc(div_u_norms)
print(h1_eocs_u)
print(l2_eocs_u)
print(l2_eocs_p)
print(div_u_eocs)


# In[17]:


# Do a pretty print of the tables using panda
import pandas as pd
from IPython.display import display

table = pd.DataFrame({'L2 error u': l2_errors_u, 'L2 EOC u ': l2_eocs_u,
                      'L2 error p': l2_errors_p, 'L2 EOC p ': l2_eocs_p})

#print(table)
display(table)


# In[ ]:


Draw(u_h, mesh, "u")


# In[ ]:


Draw(p_h, mesh, "p")


# In[ ]:


Draw(Norm(u_h), mesh, "|u|")


# In[ ]:


Draw(div(u_h), mesh, "|divu|")