from firedrake import *
import numpy as np
from firedrake.pyplot import plot
import matplotlib.pyplot as plt


mesh = UnitIntervalMesh(100)
x = SpatialCoordinate(mesh)[0]

Th = FunctionSpace(mesh, "Lagrange", 1) 
Wh = FunctionSpace(mesh, "Lagrange", 2)

Kh = MixedFunctionSpace([Wh,Wh])
Vh = MixedFunctionSpace([Th,Wh])

U = Function(Vh)
V = TestFunction(Vh)

theta,w = split(U)
eta,v = split(V)

# Material properties
t = Constant(0.0001)
kappa = 4./t**2             ## alpha is included here

f = 8*pi**3*cos(2*pi*x)
psi = Constant(-0.03)

#obstacle = Function(V, name="Obstacle")
#obstacle.interpolate(conditional(lt(psi, 0), 0, psi))

obstacle = Function(Vh.sub(1)).interpolate((Constant(-0.03)))

k = grad(theta)
gamma = w.dx(0)-theta


# Energy densities
psi_b = inner(grad(theta),grad(theta))*dx       # Bending
psi_s = inner(gamma,gamma)*dx                   # Shear

# External work
W_ext = inner(f, w)*dx

# The total energy:
J = psi_b + kappa*psi_s  - W_ext        # Energy
F = derivative(J,U,V)

bcs = [DirichletBC(Vh.sub(0), 0, "on_boundary"),
       DirichletBC(Vh.sub(1), 0, "on_boundary"),
    ]




sp = {"snes_type": "vinewtonrsls",
      "snes_monitor": None}
problem = NonlinearVariationalProblem(F, U, bcs)
solver = NonlinearVariationalSolver(problem, solver_parameters=sp)

upper = Function(Vh.sub(1)).interpolate(Constant(1e10))
solver.solve(bounds=(obstacle, upper))


(u,r) = U.subfunctions

File("output/Timoshenko-daviga.pvd").write(r)

firedrake.pyplot.plot(u,color="r") 
plt.show()