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


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

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

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.001)
kappa = 4./t**2             ## alpha is included here

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

# Create the full function in mixed space
obs = Function(Vh, name="Obstacle")
upp = Function(Vh, name="Upper")

obs_sub1 = obs.sub(1)
upp_sub1 = upp.sub(1)

# Interpolate into the subfunctions
obs_sub1.interpolate(conditional(lt(psi, 0), 0., psi))
upp_sub1.interpolate(Constant(1e3))



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


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

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

# The total energy:
J = bending + kappa*Shear  - 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)

# Solve with bounds
solver.solve(bounds=(obs, upp))


(theta,w) = U.subfunctions

File("output/Timoshenko-obs.pvd").write(theta)

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