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

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)


V = VectorFunctionSpace(mesh, "CG", 2)
bc = [DirichletBC(V.sub(1), 0, facets, 1), DirichletBC(V.sub(0), 0, facets, 3)]
n = FacetNormal(mesh)

for hypothesis in ["plane_strain", "axisymmetric"]:
    u  = Function(V, name="Displacement")
    if hypothesis == "plane_strain":
        print("Expansion of a thick cylinder in 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])
    print("---------------------------------------------")
    loading = Expression("-q*t", q=q_lim, t=0, degree=2)

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

    p = problem.get_state_variable("EquivalentPlasticStrain")
    assert (ufl.shape(p) == ())
    epsel = problem.get_state_variable("ElasticStrain")
    assert (ufl.shape(epsel)==(4, ))

    file_results = XDMFFile("results/plasticity_{}_results.xdmf".format(hypothesis))
    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(p, 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()