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

hypothesis = "axisymmetric" # axisymmetric

Re, Ri = 1.3, 1.   # external/internal radius
# elastic parameters
E = 70e3
nu = 0.3
# yield strength
sig0 = 250.
Et = E/100.
# hardening slope
H = E*Et/(E-Et)

mesh = Mesh("../meshes/thick_cylinder.xml")

facets = MeshFunction("size_t", mesh, "../meshes/thick_cylinder_facet_region.xml")
ds = Measure('ds', subdomain_data=facets)

if hypothesis == "plane_strain":
    q_lim = float(2/sqrt(3)*ln(Re/Ri)*sig0)
    measure = 1
elif hypothesis == "axisymmetric":
    x = SpatialCoordinate(mesh)
    q_lim = float(2*ln(Re/Ri)*sig0)
    measure = 2*pi*abs(x[0])

V = VectorFunctionSpace(mesh, "CG", 2)
u  = Function(V, name="Displacement")
bc = [DirichletBC(V.sub(1), 0, facets, 1), DirichletBC(V.sub(0), 0, facets, 3)]
n = FacetNormal(mesh)
loading = Expression("-q*t", q=q_lim, t=0, degree=2)

mat_prop = {"YoungModulus": E,
           "PoissonRatio": nu,
           "HardeningSlope": H,
           "YieldStrength": sig0}
material = mf.NonlinearMaterial('materials/src/libBehaviour.so', 'IsotropicLinearHardeningPlasticity',
                             hypothesis=hypothesis,
                             material_properties=mat_prop)
problem = mf.MFrontNonlinearProblem(u, material, quadrature_degree=4)
problem.set_loading(loading*dot(n, u)*measure*ds(4))
problem.bc = bc
p = problem.register_state_variable(name="EquivalentPlasticStrain")

epsel = problem.register_state_variable(name="ElasticStrain", shape=4)
print(problem.state_variables)

file_results = XDMFFile("results/plasticity_results.xdmf")
file_results.parameters["flush_output"] = True
file_results.parameters["functions_share_mesh"] = True
P0 = FunctionSpace(mesh, "DG", 0)
p_avg = Function(P0, name="Plastic_strain")

Nincr = 40
load_steps = np.linspace(0, 1.1, Nincr+1)[1:]**0.5
results = np.zeros((Nincr+1, 2))
for (i, t) in enumerate(load_steps):
    loading.t = t
    problem.solve(u.vector())

    file_results.write(u, t)

    p_avg.assign(project(epsel[0], P0))
    file_results.write(p_avg, t)
    results[i+1, :] = (u(Ri, 0)[0], t)


import matplotlib.pyplot as plt
plt.plot(results[:, 0], results[:, 1], "-o")
plt.xlabel("Displacement of inner boundary")
plt.ylabel(r"Applied pressure $q/q_{lim}$")
plt.show()