from dolfin import *
import mfront_wrapper as mf
import numpy as np
import ufl

length, width , height = 1., 40e-3, 100e-3
nb_elt_p, nb_elt_l = 10, 30
mesh = BoxMesh(Point(0, -width/2, -height/2.), Point(length, width/2, height/2.), nb_elt_l, 2, nb_elt_p)

V = VectorFunctionSpace(mesh, "CG", 2)
u = Function(V, name="Displacement")

def left(x, on_boundary):
    return near(x[0], 0) and on_boundary
def right(x, on_boundary):
    return near(x[0], length) and on_boundary

bc = [DirichletBC(V, Constant((0.,)*3), left),
      DirichletBC(V.sub(0), Constant(0.), right)]
facets = MeshFunction("size_t", mesh, 2)
AutoSubDomain(right).mark(facets, 1)
ds = Measure("ds", subdomain_data=facets)

# Building function with vx=1 on left boundary to compute reaction
v = Function(V)
bcv = DirichletBC(V.sub(0), Constant(1.), left)
bcv.apply(v.vector())

Vpost = FunctionSpace(mesh, "CG", 1)
file_results = XDMFFile("results/beam_GL_plasticity.xdmf")
file_results.parameters["flush_output"] = True
file_results.parameters["functions_share_mesh"] = True


selfweight = Expression(("0", "0", "-t*qmax"), t=0., qmax = 50e6, degree=0)
material = mf.MFrontNonlinearMaterial('../materials/src/libBehaviour.so',
                                      "LogarithmicStrainPlasticity")
problem = mf.MFrontNonlinearProblem(u, material)
problem.set_loading(dot(selfweight, u)*dx)
problem.bc = bc

p = problem.get_state_variable(name="EquivalentPlasticStrain")
assert (ufl.shape(p) == ())
epsel = problem.get_state_variable(name="ElasticStrain")
print(ufl.shape(epsel))

prm = problem.solver.parameters
prm["absolute_tolerance"] = 1e-6
prm["relative_tolerance"] = 1e-6

P0 = FunctionSpace(mesh, "DG", 0)
p_avg = Function(P0, name="Plastic strain")
Nitermax, tol = 20, 1e-4  # parameters of the Newton-Raphson procedure
Nincr = 30
load_steps = np.linspace(0., 1., Nincr+1)
file_results.write(u, 0)
file_results.write(p_avg, 0)
results = np.zeros((Nincr+1, 3))
for (i, t) in enumerate(load_steps[1:]):
    selfweight.t = t

    problem.solve(u.vector())
    results[i+1, 0] = assemble(-u[2]*ds(1))/(width*height)
    results[i+1, 1] = t
    results[i+1, 2] = assemble(action(-problem.L, v))

    file_results.write(u, t)
#    file_results.write(p_avg, t)

import matplotlib.pyplot as plt
plt.figure()
plt.plot(results[:, 0], results[:, 1], "-o")
plt.xlabel(r"Displacement")
plt.ylabel("Load")
plt.figure()
plt.plot(results[:, 1], results[:, 2], "-o")
plt.xlabel(r"Load")
plt.ylabel("Horizontal Reaction")
plt.show()