from dolfin import *
from fenicstools import Probes
from numpy import linspace, repeat, where, resize
import os

parameters["form_compiler"]["optimize"] = True
parameters["form_compiler"]["cpp_optimize"] = True
parameters["form_compiler"]["representation"] = "quadrature"
parameters["form_compiler"]["cpp_optimize_flags"] = "-O3 --fast-math"

if not os.path.isfile("cylinder.xml"):
    os.system("wget https://raw.githubusercontent.com/mikaem/Oasis/master/mesh/cylinder.geo -O cylinder.geo")
    os.system("gmsh cylinder.geo -2 -o cylinder.msh")
    os.system("dolfin-convert cylinder.msh cylinder.xml")
    os.system("rm cylinder.msh")    

mesh = Mesh("cylinder.xml")
H = 0.41
L = 2.2
D = 0.1
center = 0.2

Um = 0.3
Umean = 2./3.*Um
Re = 20.0
nu = 2./3.*Um*D/Re

inlet = Expression(("4.*{0}*x[1]*({1}-x[1])/pow({1}, 2)".format(Um, H), "0"))
ux = Expression("-1.5*x[1]")
uy = Expression("1.5*(x[0]-{})".format(center))

V = VectorFunctionSpace(mesh, 'CG', 2)
#V = FunctionSpace(mesh, 'CG', 1)
#V = VectorFunctionSpace(mesh, 'CR', 1)
#B = VectorFunctionSpace(mesh, "Bubble", 3)
#V = V + B

Q = FunctionSpace(mesh, 'CG', 1)
VQ = V * Q
u, p = TrialFunctions(VQ)
v, q = TestFunctions(VQ)
ur = Expression(("-1.5*(x[1]-{0})".format(center), "1.5*(x[0]-{0})".format(center)))
up_ = Function(VQ)
u_, p_ = split(up_)
n = FacetNormal(mesh)

F = inner(dot(grad(u_), u_), v)*dx + nu*inner(grad(u_), grad(v))*dx \
    - inner(p_, div(v))*dx - inner(q, div(u_))*dx

# Specify boundary conditions
Inlet = AutoSubDomain(lambda x, on_bnd: on_bnd and x[0] < 1e-8)
Wall = AutoSubDomain(lambda x, on_bnd: on_bnd and near(x[1]*(H-x[1]), 0))
Cyl = AutoSubDomain(lambda x, on_bnd: on_bnd and x[0]>1e-6 and x[0]<1 and x[1] < 3*H/4 and x[1] > H/4)
Outlet = AutoSubDomain(lambda x, on_bnd: on_bnd and x[0] > L-1e-8)

bc1 = DirichletBC(VQ.sub(0), Constant((0, 0)), Wall)
bc2 = DirichletBC(VQ.sub(0), ur, Cyl)
bc3 = DirichletBC(VQ.sub(0), inlet, Inlet)

solve(F == 0, up_, bcs=[bc1, bc2, bc3], 
      solver_parameters={'newton_solver':{'maximum_iterations': 10, 
                                          'error_on_nonconvergence': False,
                                          'linear_solver': 'mumps'}})

plot(u_)

R = VectorFunctionSpace(mesh, 'R', 0)
c = TestFunction(R)

tau = -p_*Identity(2)+nu*(grad(u_)+grad(u_).T)

ff = FacetFunction("size_t", mesh, 0)
Cyl.mark(ff, 1)
n = FacetNormal(mesh)
ds = ds(subdomain_data=ff)

forces = assemble(dot(dot(tau, n), c)*ds(1)).array()*2/Umean**2/D

print "Cd = {}, CL = {}".format(*forces)

xx = linspace(0, L, 10000)
x = resize(repeat(xx, 2), (10000, 2))
x[:, 1] = 0.2
probes = Probes(x.flatten(), VQ)
probes(up_)
nmax = where(probes.array()[:, 0] < 0)[0][-1]
print "L = ", x[nmax, 0]-0.25
print "dP = ", up_(Point(0.15, 0.2))[2] - up_(Point(0.25, 0.2))[2]
print "Global divergence error ", assemble(dot(u_, n)*ds()), assemble(div(u_)*div(u_)*dx())