#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Feb  1 08:49:11 2019

@author: bleyerj
"""
from dolfin import *
import mfront_wrapper as mf
import numpy as np
from ufl import diag

width = 1.
height = 1.
mesh = RectangleMesh(Point(0., 0.), Point(width, height), 20, 10)

Ve = VectorElement("CG", mesh.ufl_cell(), 1)
V = FunctionSpace(mesh, MixedElement([Ve, Ve]))
u = Function(V, name="Displacements")
(u1, u2) = split(u)

def bottom(x, on_boundary):
    return near(x[1], 0) and on_boundary
def top(x, on_boundary):
    return near(x[1], height) and on_boundary
def left(x, on_boundary):
    return near(x[0], 0) and on_boundary


bc = [DirichletBC(V.sub(0).sub(0), Constant(0), left),
      DirichletBC(V.sub(1).sub(0), Constant(0), left),
      DirichletBC(V.sub(0).sub(1), Constant(0), bottom),
      DirichletBC(V.sub(1).sub(1), Constant(0), bottom),
      DirichletBC(V.sub(0).sub(1), Constant(-1), top),
      DirichletBC(V.sub(1).sub(1), Constant(-1), top)]

facets = MeshFunction("size_t", mesh, 1)
ds = Measure("ds", subdomain_data=facets)


material = mf.MFrontNonlinearMaterial("materials/src/libBehaviour.so",
                                      "TwoPhasesGeneralizedElasticity",
                                      hypothesis="plane_strain")

problem = mf.MFrontNonlinearProblem(u, material, quadrature_degree=0)
problem.bc = bc
problem.register_gradient("MatrixStrain", sym(grad(u1)), symmetric=True)
problem.register_gradient("InclusionStrain", sym(grad(u2)), symmetric=True)
problem.register_gradient("RelativeDisplacement", diag(u2-u1), symmetric=True)

problem.solve(u.vector())

u1 = u.sub(0, True)
u2 = u.sub(1, True)
x = np.linspace(0, width, 100)
import matplotlib.pyplot as plt
plt.figure()
plt.plot(x, np.array([u1(xi, height/2)[0] for xi in x]), label=r"$u_x$ (matrix)")
plt.plot(x, np.array([u2(xi, height/2)[0] for xi in x]), label=r"$u_x$ (inclusion)")
plt.legend()
plt.show()