from dolfin import *

mesh = Mesh("cyl_Z.xml")
#mesh = refine(mesh)

V = VectorFunctionSpace(mesh, 'CG', 1)
Q = FunctionSpace(mesh, 'CG', 1)
S = TensorFunctionSpace(mesh, 'CG', 1)
u = TrialFunction(V)
v = TestFunction(V)
n = FacetNormal(mesh)

# Lag SubDomains av grenseflatene i topp og bunn
right = AutoSubDomain(lambda x, on_bnd: near(x[0], 10.))
left = AutoSubDomain(lambda x, on_bnd: near(x[0], 0.))

# Marker grenseflatene med en FacetFunction 
mf = FacetFunction('uint', mesh)
mf.set_all(0)
left.mark(mf, 1)
right.mark(mf, 2)

# Sett konstanterq
Lambda = Constant(20)
mu = Constant(10)
L = Constant(10)
q = Constant(0.01) # Vridningsvinkel / L

# Lag en torsjon grensebetingelse for P*n
traction = Expression(("0", "-q*mu*x[2]", "q*mu*x[1]"), element=V.ufl_element(), q=q(0), mu=mu(0))

# Lag en grensebetingelse for ingen forskyvning til venstre i cylinderen
bc1 = DirichletBC(V, Constant((0, 0, 0)), left)
bc2 = DirichletBC(V, Expression(("0", "-q*L*x[2]", "q*L*x[1]"), element=V.ufl_element(), q=q(0), L=L(0)), right)

# Spenningstensor
P = Lambda * div(u) * Identity(3) + mu * (grad(u) + grad(u).T)

# Likevektslikningen
F = inner(grad(v), P)*dx - inner(v, traction)*ds(2)

# Assemble and solve problem
A = assemble(lhs(F), exterior_facet_domains=mf)
b = assemble(rhs(F), exterior_facet_domains=mf)
u = Function(V)
bc1.apply(A, b)
bc2.apply(A, b)
solve(A, u.vector(), b)

# Legg resultatet i fil (kan leses av paraview)
f = File('u_torsjon.pvd')
f << u
#plot(u[0], title="Forskyvning i x-retning")
plot(u[1], title="Forskyvning i y-retning")
plot(u[2], title="Forskyvning i z-retning")

# Beregn resultantmomentet av spenningen paa endeflaten
r = interpolate(Expression("sqrt(x[2]*x[2]+ x[1]*x[1])"), Q)
x, y, z = interpolate(Expression(("x[0]", "x[1]", "x[2]")), V).split()
ii = as_vector((x, y/r, z/r))
P = Lambda * div(u) * Identity(3) + mu * (grad(u) + grad(u).T) # Trenger P som Funksjon, ikke argument
moment = cross(r*ii, dot(P, n))

print "Resultantmoment for (x = L): "
print "Analytisk : ", pi / 2 * mu(0) * q(0)
print "Numerisk  : ", assemble(dot(moment, n)*ds(2), exterior_facet_domains=mf)

interactive()

Pn = dot(P, n)
Pnn = dot(Pn, n)
Pt = Pn - Pnn*n
Pnt = sqrt(dot(cross(Pn, n), cross(Pn, n)))
tt = Pt / Pnt