import matplotlib.pyplot as plt
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)]

file_results = XDMFFile("results/finie_strain_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, bcs=bc)
problem.set_loading(dot(selfweight, u)*dx)

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

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

P0 = FunctionSpace(mesh, "DG", 0)
p_avg = Function(P0, name="Plastic strain")

Nincr = 30
load_steps = np.linspace(0., 1., Nincr+1)
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] = u(length, 0, 0)
    results[i+1, 1] = t

    file_results.write(u, t)
    p_avg.assign(project(p, V0))
    file_results.write(p_avg, t)

plt.figure()
plt.plot(results[:, 0], results[:, 1], "-o")
plt.xlabel(r"Displacement")
plt.ylabel("Load")