from ngsolve import *

# viscosity
nu = 0.001

# timestepping parameters
tau = 0.001
tend = 3
opening = 0.1
opStart = 1.41
opEnd = opStart + opening
from netgen.geom2d import SplineGeometry
geo = SplineGeometry()
#pnts =[
p1 = geo.AppendPoint(0,0)
p2 = geo.AppendPoint(2,0)
p3 = geo.AppendPoint(2,0.41)
p4 = geo.AppendPoint(opEnd,0.41)
p5 = geo.AppendPoint(opEnd+opening,0.51)
p6 = geo.AppendPoint(opEnd,0.61)
p7 = geo.AppendPoint(opStart,0.61)
p8 = geo.AppendPoint(opStart-opening,0.51)
p9 = geo.AppendPoint(opStart,0.41)
p10 = geo.AppendPoint(0,0.41)


geo.Append(["line", p1,p2],bc="wall")
geo.Append(["line", p2,p3],bc="outlet")
geo.Append(["line", p3,p4],bc="wall")
geo.Append(["line", p4,p5],bc="wall")
geo.Append(["line", p5,p6],bc="wall")
geo.Append(["line", p6,p7],bc="wall")
geo.Append(["line", p7,p8],bc="wall")
geo.Append(["line", p8,p9],bc="wall")
geo.Append(["line", p9,p10],bc="wall")
geo.Append(["line", p10,p1],bc="inlet")


#geo.AddRectangle( (0, 0), (2, 0.41), bcs = ("wall", "outlet", "wall", "inlet"))
#geo.AddRectangle( (1, 0.41), (1.41, 0.82), bcs = ("inlet", "wall", "wall", "wall"))
geo.AddCircle ( (0.2, 0.2), r=0.05, leftdomain=0, rightdomain=1, bc="cyl", maxh=0.02)

mesh = Mesh( geo.GenerateMesh(maxh=0.07))

mesh.Curve(3)

V = VectorH1(mesh,order=3, dirichlet="wall|cyl|inlet")
Q = H1(mesh,order=2)

X = V*Q

u,p = X.TrialFunction()
v,q = X.TestFunction()

stokes = nu*InnerProduct(grad(u), grad(v))+div(u)*q+div(v)*p - 1e-12*p*q
a = BilinearForm(X, symmetric=True)
a += stokes*dx
a.Assemble()

# nothing here ...
f = LinearForm(X)   
f.Assemble()

# gridfunction for the solution
gfu = GridFunction(X)

# parabolic inflow at inlet:
uin = CoefficientFunction( (1.5*4*y*(0.41-y)/(0.41*0.41), 0) )
gfu.components[0].Set(uin, definedon=mesh.Boundaries("inlet"))

# solve Stokes problem for initial conditions:
inv_stokes = a.mat.Inverse(X.FreeDofs())

res = f.vec.CreateVector()
res.data = f.vec - a.mat*gfu.vec
gfu.vec.data += inv_stokes * res


# matrix for implicit Euler 
mstar = BilinearForm(X, symmetric=True)
mstar += (u*v + tau*stokes)*dx
mstar.Assemble()
inv = mstar.mat.Inverse(X.FreeDofs(), inverse="sparsecholesky")

# the non-linear term 
conv = BilinearForm(X, nonassemble = True)
conv += (grad(u) * u) * v * dx

# for visualization
Draw (Norm(gfu.components[0]), mesh, "velocity", sd=3)

# implicit Euler/explicit Euler splitting method:
t = 0
with TaskManager():
    while t < tend:
        print ("t=", t, end="\r")

        conv.Apply (gfu.vec, res)
        res.data += a.mat*gfu.vec
        gfu.vec.data -= tau * inv * res    

        t = t + tau
        Redraw()