from dolfin import *

# Mesh Params ################
E, nu = 1.0, 0.4 #material params
mesh = UnitSquareMesh(100, 100)

print 'Creating QUADRATIC Lagrangian function space...'
Vv = VectorFunctionSpace(mesh, "Lagrange", 2)
Vf = FunctionSpace(mesh, "Lagrange", 2)
Vt = TensorFunctionSpace(mesh, "Lagrange", 2)

# P is stress vector of stress values normal and tangent to boundary
P = Expression(('(x[0]-xc)','(x[0]-xc)'),xc=0.5)

class LeftBoundary(SubDomain):
    def inside(self, x, on_boundary):
        tol = 1E-7
        return on_boundary and abs(x[0]) < tol

class RightBoundary(SubDomain):
    def inside(self, x, on_boundary):
        tol = 1E-7
        return on_boundary and abs(x[0] - 1) < tol

left_boundary = LeftBoundary()
right_boundary = RightBoundary()

boundary_parts = MeshFunction("size_t", mesh, mesh.topology().dim() - 1)
boundary_parts.set_all(0)
right_boundary.mark(boundary_parts, 1)
left_boundary.mark(boundary_parts, 2)
ds = Measure("ds")[boundary_parts]

##############################
# Define variational problem #
##############################
print 'Defining test and trial functions...'
v = TestFunction(Vf)
tau = TestFunction(Vv)
phi = TrialFunction(Vf)
u = TrialFunction(Vv)

# functions for variational problem
d = u.geometric_dimension()
epsilon = sym(nabla_grad(u))
# PLANE STRESS
sigma = E/(1+nu)* (epsilon) + E*nu/((1-nu**2)) * tr(epsilon)*Identity(d)

# Define normal and tangent directions on mesh
norm = FacetNormal(mesh)
tang = as_vector([norm[1], -norm[0]])

tau_n = dot(tau,norm)
tau_t = dot(tau,tang)
P_n = dot(P,norm)
P_t = dot(P,tang)

#################
# SOLVE
#################
au = inner(sigma, sym(nabla_grad(tau)))*dx
Lu = dot(P_n,tau_n)*ds(1) + dot(P_n,tau_n)*ds(2) + dot(P_t,tau_t)*ds(1) + dot(P_t,tau_t)*ds(2)
u = Function(Vv)
solve(au == Lu, u)

## View the results
sigma = E/(1+nu)* (sym(nabla_grad(u))) + E*nu/((1-nu**2)) * tr(sym(nabla_grad(u)))*Identity(d)
s_V = project( sigma, Vt)
sxx = project(s_V[0,0], Vf);
sxy = project(s_V[0,1], Vf);
syx = project(s_V[1,0], Vf);
syy = project(s_V[1,1], Vf);
plot(sxx, interactive=True, rescale=False,scalarbar=True)
plot(syy, interactive=True, rescale=False,scalarbar=True)
plot(sxy, interactive=True, rescale=False,scalarbar=True)