#!/usr/bin/env python
# coding: utf-8

# $\newcommand{\sym}{\operatorname{sym}}
# \newcommand{\skew}{\operatorname{skew}}
# \renewcommand{\tr}{\operatorname{tr}}
# \newcommand{\T}{^\text{T}}
# \newcommand{\utilde}[1]{\tilde{#1}}
# \newcommand{\uttilde}[1]{\tilde{\tilde{#1}}}
# \newcommand{\tgrad}{\utilde{\nabla}}
# \newcommand{\bX}{\boldsymbol{X}}
# \newcommand{\bsig}{\boldsymbol{\sigma}}
# \newcommand{\btheta}{\boldsymbol{\theta}}
# \newcommand{\bvarepsilon}{\boldsymbol{\varepsilon}}
# \newcommand{\bepsilon}{\boldsymbol{\epsilon}}
# \newcommand{\bgamma}{\boldsymbol{\gamma}}
# \newcommand{\bchi}{\boldsymbol{\kappa}}
# \newcommand{\bX}{\boldsymbol{X}}
# \newcommand{\bxi}{\boldsymbol{\xi}}
# \newcommand{\ba}{\boldsymbol{a}}
# \newcommand{\be}{\boldsymbol{e}}
# \newcommand{\bt}{\boldsymbol{t}}
# \newcommand{\bu}{\boldsymbol{u}}
# \newcommand{\bv}{\boldsymbol{v}}
# \newcommand{\bx}{\boldsymbol{x}}
# \newcommand{\bF}{\boldsymbol{F}}
# \newcommand{\bI}{\boldsymbol{I}}
# \newcommand{\bM}{\boldsymbol{M}}
# \newcommand{\bN}{\boldsymbol{N}}
# \newcommand{\bQ}{\boldsymbol{Q}}
# \newcommand{\bR}{\boldsymbol{R}}
# \newcommand{\bV}{\boldsymbol{V}}
# \newcommand{\be}{\boldsymbol{e}}$

# # Linear shell model
# 
# This tour is dedicated to the implementation of small-strain elastic shell model. In the [fenics-shells project](https://fenics-shells.readthedocs.io), the shell surface is described by a mapping between a parametric 2D domain and the shell initial stress-free configuration. This approach makes it impossible to tackle non-manifold surfaces or closed surfaces for which the mapping would exhibit singularities.
# 
# In the present implementation, we will instead work directly with the shell surface (either manifold or non-manifold) described as an assembly of planar facets. In this respect, the approach will bear important similarities with the [Elastic 3D beam structures](https://comet-fenics.readthedocs.io/en/latest/demo/beams_3D/beams_3D.html). In particular, we will take advantage of `UFL` ability to compute gradients over a manifold.
# 
# Regarding the interpolation spaces, we will use a `P2/CR` mixed space for displacements and rotations, i.e. the same choice as in the [Locking-free Reissner-Mindlin plate with Crouzeix-Raviart interpolation](https://comet-fenics.readthedocs.io/en/latest/demo/reissner_mindlin_crouzeix_raviart/reissner_mindlin_CR.html) tour. The present shell model will therefore be a small-displacement version of the nonlinear shell model developed in [[CAM03]](#References).
# 
# > The implementation of the fully nonlinear shell model of [[CAM03]](#References) is developed in a separate project for research purposes. You can directly contact the author for more details on this matter.
# 
# For the mesh generation, we refer to the [Generating a shell model for a shallow arch I-beam with the Gmsh Python API](I_beam_gmsh.ipynb) tour.

# ## Shell kinematics and 3D linearized strain
# 
# The shell initial reference configuration consists of piecewise flat portions and $\bxi$ in $\mathbb{R}^3$ will denote an initial point on this surface. $(\be_1,\be_2,\be_3)$ will be a reference local frame with $\be_3$ the normal to the shell surface, this frame being orthonormal and piecewise constant. A material point $\bX$ in the shell reference configuration will then be given by:
# \begin{equation}
# \bX(\bxi,\zeta) = \bxi + \zeta \be_3 \quad \text{with } \zeta \in [-h/2;h/2]
# \end{equation}
# where $h$ is the shell thickness
# 
# The shell kinematics will be described by its mid-surface displacement $\bu$ and the infinitesimal rotation vector $\btheta$ of its director. The new normal director is $\ba_3 = \bR(\btheta)\be_3 = \be_3+\btheta\times\be_3$ with $\bR$ the infinitesimal rotation matrix associated with $\btheta$. Neglecting any thickness change in the shell kinematics, the material point $\bx$ in the deformed configuration associated with $\bX$ will then be given by:
# 
# \begin{equation}
# \bx = \bxi + \bu(\bxi) + \zeta\ba_3 = \bxi + \bu(\bxi) + \zeta(\be_3+\btheta(\bxi)\times\be_3)
# \end{equation}
# Differentiating with respect to $\bX$, we get:
# 
# \begin{align}
# d\bx &= d\bxi + \nabla\bu\cdot d\bxi + \zeta(\nabla\btheta\cdot d\bxi)\times\be_3 + d\zeta(\be_3+\btheta\times\be_3) \\
# &= d\bX + \nabla\bu\cdot d\bxi - \zeta (\be_3\times \nabla\btheta )\cdot d\bxi - d\zeta (\be_3\times\btheta) + \text{ h.o.t.} \notag
# \end{align}
# 
# where we retained only up to first order terms in $\bu,\btheta$. The 3D deformation gradient is then given by:
# \begin{equation}
# \bF = \dfrac{d\bx}{d\bX} = \bI + \tgrad\bu  - \zeta \be_3\times \tgrad\btheta  - (\be_3\times\btheta)\otimes\be_3 
# \end{equation} 
# 
# where we introduced the in-plane gradient  (i.e. the gradient with respect to the shell local tangent plane $(\be_1,\be_2)$) as follows $\tgrad\bv = \partial_1\bv \otimes \be_1+ \partial_2\bv \otimes \be_2$. More generally, we will use the following notation $\utilde{\bv} = v_1\be_1+v_2\be_2$ to denote the in-plane part of $\bv$.
# 
# The linearized strain tensor is then given by:
# \begin{equation}
# \bvarepsilon = \sym\left(\tgrad\bu - \zeta \be_3\times \tgrad\btheta  - (\be_3\times\btheta)\otimes\be_3 \right) +  \epsilon(\zeta) \be_3\otimes\be_3
# \end{equation} 
# where $\sym$ denotes the symmetrized part of a tensor and where we added an incompatible out-of-plane strain $\epsilon(\zeta)$ which will enable to enforce a plane-stress constraint. The 3D strain can be split into its in-plane components $\utilde{\utilde{\bvarepsilon}}$, out-of-plane shear components $\bgamma = 2\utilde{\bvarepsilon}_3$ and out-of-plane transverse components $\varepsilon_{33}$:
# \begin{align}
# \uttilde{\bvarepsilon} &= \sym\left(\tgrad\utilde{\bu} - \zeta \be_3\times \tgrad\btheta \right)= \boldsymbol{\epsilon} - \zeta \bchi \\
# \bgamma &= \tgrad u_3 - \be_3\times\btheta \tag{1}\\
# \varepsilon_{33} &= \epsilon(\zeta)
# \end{align}
# where $\bepsilon=\sym(\tgrad\utilde{\bu})$ is the membrane strain and $\bchi = \sym(\be_3\times \tgrad\btheta)$ the bending curvature. We see that we recover the definition of the curvature and shear strain of the Reissner-Mindlin plate model.

# ## Elastic energy density
# 
# 
# The internal work of deformation density per shell unit surface is then given by:
# 
# \begin{align}
# w_\text{def} &= \int_{-h/2}^{h/2} \bsig:\bvarepsilon d\zeta \\
# &= \left(\int_{-h/2}^{h/2} \uttilde{\bsig} d\zeta\right):\bepsilon + \left(\int_{-h/2}^{h/2} (-\zeta)\uttilde{\bsig} d\zeta\right):\bchi \notag\\
# &\phantom{=}+  \left(\int_{-h/2}^{h/2} \utilde{\bsig}_3 d\zeta\right)\cdot\bgamma + \int_{-h/2}^{h/2} \sigma_{33}\epsilon(\zeta)d\zeta \notag \\
# &= \bN:\bepsilon + \bM:\bchi + \bQ\cdot\bgamma + \int_{-h/2}^{h/2} \sigma_{33}\epsilon(\zeta)d\zeta \notag
# \end{align}
# 
# where $\bN$ is the shell membrane tensor, $\bM$ the bending moment tensor and $\bQ$ the shear force vector appearing in duality with $\bepsilon,\bchi,\bgamma$ respectively. The out-of-plane stress $\sigma_{33}$ appears in duality with the out-of-plane strain $\epsilon(\zeta)$. The latter is in fact a purely local variable which can be locally eliminated to enforce the plane stress condition $\sigma_{33}=0$.
# 
# Indeed, the shell free energy density can be defined as:
# 
# \begin{align}
# \Psi &= \int_{-h/2}^{h/2} \psi(\bvarepsilon) d\zeta \\
# &= \int_{-h/2}^{h/2} \frac{1}{2}(\lambda (\tr\bvarepsilon)^2 + 2\mu \bvarepsilon:\bvarepsilon) d\zeta
# \end{align}
# 
# for isotropic linear elasticity.
# 
# When minimizing this energy over the shell degrees of freedom, the minimization over $\epsilon$ is in fact local and we can replace the above free-energy by its plane-stress counterpart:
# 
# 
# \begin{equation}
# \begin{array}{rl}\Psi_\text{ps} &= \displaystyle{\int_{-h/2}^{h/2} \inf_{\epsilon(z)}\psi(\bvarepsilon) d\zeta}\\
# & =\displaystyle{\int_{-h/2}^{h/2}\left(\psi_\text{ps}(\uttilde{\bvarepsilon}) + \dfrac{1}{2}\mu \bgamma^2\right) d\zeta  }
# \end{array}
# \end{equation}
# 
# where $\psi_\text{ps}$ is the elastic strain energy density for a 2D plane-stress material:
# 
# \begin{equation}
# \psi_\text{ps}(\uttilde{\bvarepsilon}) = \frac{1}{2}\left(\lambda_\text{ps} (\tr\uttilde{\bvarepsilon})^2 + 2\mu \uttilde{\bvarepsilon}:\uttilde{\bvarepsilon}\right) 
# \end{equation}
# 
# with the plane-stress pseudo Lamé coefficient (see the [2D linear elasticity](https://comet-fenics.readthedocs.io/en/latest/demo/elasticity/2D_elasticity.py.html) tour):
# 
# \begin{equation}
# \lambda_\text{ps} = \dfrac{2\lambda\mu}{\lambda+2\mu}
# \end{equation}
# 
# Taking the derivative of the strain energy with respect to the 3D strain components gives the local 3D stress-strain constitutive relation:
# 
# \begin{align}
# \uttilde{\bsig} &= \lambda_\text{ps} \tr(\uttilde{\bvarepsilon})\mathbf{1} + 2\mu \bvarepsilon = \mathbb{C}_\text{ps}:\uttilde{\bvarepsilon} \\
# \utilde{\bsig}_{3} &= \mu \bgamma \\
# \sigma_{33} &= 0
# \end{align}
# 
# where $\mathbb{C}_\text{ps}$ is the plane-stress elasticity tensor.

# ## Stress-resultant constitutive equations
# 
# Plugging relations (1) into the previous 3D constitutive relation yields the following stress-resultant equations, assuming a homogeneous shell across the thickness:
# 
# \begin{align}
# \bN &= \int_{-h/2}^{h/2} \uttilde{\bsig} d\zeta = h \int_{-h/2}^{h/2} \mathbb{C}_\text{ps}:(\boldsymbol{\epsilon} - \zeta \bchi) d\zeta = h\mathbb{C}_\text{ps}:\boldsymbol{\epsilon} \\
# \bM &= \int_{-h/2}^{h/2} (-\zeta)\uttilde{\bsig} d\zeta = h \int_{-h/2}^{h/2} (-\zeta)\mathbb{C}_\text{ps}:(\boldsymbol{\epsilon} - \zeta \bchi) d\zeta = \dfrac{h^3}{12}\mathbb{C}_\text{ps}:\boldsymbol{\bchi}\\
# \bQ &= \int_{-h/2}^{h/2} \utilde{\bsig}_3 d \zeta = \mu h \bgamma
# \end{align}
# 
# > Note that with the present purely kinematic approach, no shear correction factor appears for the shear force constitutive relation.
# 
# ## FEniCS implementation
# 
# ### Import statements and physical parameters
# 
# We first load the relevant packages and some UFL functions. We then define the corressponding physical parameters. In particular, we consider here a uniformly distributed loading of unit intensity along the downward global $Z$ direction and of intensity 0.2 along the $Y$ direction perpendicular to the arch plane in order to induce a combined bending and twisting deflection mode.

# In[10]:


from dolfin import *
from ufl import Jacobian, replace
import numpy as np
import meshio

# material parameters
thick = Constant(1e-3)
E = Constant(210e9)
nu = Constant(0.3)
lmbda = E * nu / (1 + nu) / (1 - 2 * nu)
mu = E / 2 / (1 + nu)
lmbda_ps = 2 * lmbda * mu / (lmbda + 2 * mu)

# loading (self-weight)
f = Constant((0, 0.2, -1))


# ### Loading the mesh and computing a local tangent frame
# 
# We now load the mesh as well as the facet MeshFunction from the XDMF files generated from a Gmsh script using Gmsh's Python API (see [this notebook](I_beam_gmsh.ipynb)).
# 
# We then define the `local_frame` function which returns the UFL representation of the local frame  $(\be_1,\be_2,\be_3)$. We first use UFL's `Jacobian` function to retrieve the Jacobian of the mapping from reference cells to spatial coordinates. Contrary to the `fenics-shells` approach which uses a global analytical mapping between the reference domain and the shell surface, we use here the reference element mapping of the finite-element method. For a shell, the Jacobian is of size $3\times 2$ with both columns $\bt_1,\bt_2$ being a vector of the local tangent plane. We therefore compute $\be_3$, the normal to the mid-surface from $\bt_1\times \bt_2$. We then have to choose a convention to compute $\be_1,\be_2$, knowing $\be_3$. Our convention is that $\be_1$ lies in the plane orthogonal to $\be_Y$ and $\be_3$. If $\be_3$ happens to be colinear to $\be_Y$, we set $\be_1=\be_Z$. $\be_2$ then follows from $\be_2=\be_3\times\be_1$.

# In[11]:


filename = "I_beam.xdmf"

mesh = Mesh()
with XDMFFile(filename) as mesh_file:
    mesh_file.read(mesh)


mvc = MeshValueCollection("size_t", mesh, 1)
with XDMFFile(filename.replace(".", "_facet_region.")) as infile:
    infile.read(mvc, "name_to_read")
facets = MeshFunction("size_t", mesh, mvc)

def local_frame(mesh):
    t = Jacobian(mesh)
    if mesh.geometric_dimension() == 2:
        t1 = as_vector([t[0, 0], t[1, 0], 0])
        t2 = as_vector([t[0, 1], t[1, 1], 0])
    else:
        t1 = as_vector([t[0, 0], t[1, 0], t[2, 0]])
        t2 = as_vector([t[0, 1], t[1, 1], t[2, 1]])
    e3 = cross(t1, t2)
    e3 /= sqrt(dot(e3, e3))
    ey = as_vector([0, 1, 0])
    ez = as_vector([0, 0, 1])
    e1 = cross(ey, e3)
    norm_e1 = sqrt(dot(e1, e1))
    e1 = conditional(lt(norm_e1, 0.5), ez, e1 / norm_e1)

    e2 = cross(e3, e1)
    e2 /= sqrt(dot(e2, e2))
    return e1, e2, e3


frame = local_frame(mesh)
e1, e2, e3 = frame

VT = VectorFunctionSpace(mesh, "DG", 0, dim=3)

ffile = XDMFFile("output.xdmf")
ffile.parameters["functions_share_mesh"] = True
for (i, ei) in enumerate(frame):
    ei = Function(VT, name="e{}".format(i + 1))
    ei.assign(project(frame[i], VT))
    ffile.write(ei, 0)


# The local frame is then stored in an output XDMFFile. Note that when it is represented via `DG0` Functions but Paraview does not seem to handle very well such fields on a surface shell. The fields seem to be displayed as `CG1` Functions, inducing erratic orientations near the boundary. Overall, the orientation is however consistent with what has been defined ($\be_1$ in blue, $\be_2$ in green, $\be_3$ in red).
# 
# <center><img src="local_frame.png" width="800"></center>
# 
# 
# ### FunctionSpace choice and strain measures
# 
# We now use the afore-mentioned `P2/CR`interpolation for the displacement $\bu$ and the rotation $\btheta$ variables using a `MixedElement`.

# In[12]:


Ue = VectorElement("Lagrange", mesh.ufl_cell(), 2, dim=3)
Te = VectorElement("CR", mesh.ufl_cell(), 1, dim=3)
V = FunctionSpace(mesh, MixedElement([Ue, Te]))

v = Function(V)
u, theta = split(v)

v_ = TestFunction(V)
u_, theta_ = split(v_)
dv = TrialFunction(V)


# We then define the in-plane projector and the in-plane (tangential) gradient operator ($\tilde\nabla$ above) as follows:

# In[13]:


def vstack(vectors):
    """Stack a list of vectors vertically."""
    return as_matrix([[v[i] for i in range(len(v))] for v in vectors])


def hstack(vectors):
    """Stack a list of vectors horizontally."""
    return vstack(vectors).T


# In-plane projection
P_plane = hstack([e1, e2])

def t_grad(u):
    """Tangential gradient operator"""
    g = grad(u)
    return dot(g, P_plane)


# We then extract gradient of the in-plane displacement $\tilde\nabla\tilde\bu$ and define the membrane strain $\bepsilon$ as its symmetric part. We similarly define the bending curvature $\bchi$ and shear strain $\bgamma$.

# In[14]:


t_gu = dot(P_plane.T, t_grad(u))
eps = sym(t_gu)
kappa = sym(dot(P_plane.T, t_grad(cross(e3, theta))))
gamma = t_grad(dot(u, e3)) - dot(P_plane.T, cross(e3, theta))

eps_ = replace(eps, {v: v_})
kappa_ = replace(kappa, {v: v_})
gamma_ = replace(gamma, {v: v_})


# ### Stress measures
# 
# We then define the corresponding stress measure using linear elastic constitutive laws.
# 
# 

# In[15]:


def plane_stress_elasticity(e):
    return lmbda_ps * tr(e) * Identity(2) + 2 * mu * e

N = thick * plane_stress_elasticity(eps)
M = thick ** 3 / 12 * plane_stress_elasticity(kappa)
Q = mu * thick * gamma


# ### Drilling rotation stabilization
# 
# A classical problem in shell models involving 6 degrees of freedom (3D rotation) is the absence of any constraint on the drilling rotation 
# $\theta_3=\btheta\cdot\be_3$. However, this degree of freedom is necessary to tackle non-smooth junctions
# between planar shell facets which have a different normal vector. In our
# implementation, we propose to add an additional quadratic energy penalizing the drilling strain, as commonly done for elastic shells.
# The drilling strain is obtained from the skew symmetric in-plane component of the transformation gradient and the drilling rotation: 
# \begin{equation}
# \varpi = \dfrac{1}{2}(u_{1,2}-u_{2,1})+\theta_3
# \end{equation}
# 
# We consider an additional drilling contribution to the work of deformation given by:
# 
# \begin{equation}
# w_\text{def,drilling} = E h^3 \varpi^2
# \end{equation}
# where the drilling stiffness $Eh^3$ is the usually recommended value in the literature.

# In[16]:


drilling_strain = (t_gu[0, 1] - t_gu[1, 0]) / 2 - dot(theta, e3)
drilling_strain_ = replace(drilling_strain, {v: v_})
drilling_stress = E * thick ** 3 * drilling_strain


# ### Boundary conditions and resolution
# We then apply fixed boundary conditions on both ends and finally form the corresponding bilinear form and linear right-hand side form before solving the corresponding system.

# In[17]:


bc = [
    DirichletBC(V.sub(0), Constant((0.0,) * 3), facets, 1),
    DirichletBC(V.sub(0), Constant((0.0,) * 3), facets, 2),
]

Wdef = (
    inner(N, eps_)
    + inner(M, kappa_)
    + dot(Q, gamma_)
    + drilling_stress * drilling_strain_
) * dx
a = derivative(Wdef, v, dv)
Wext = dot(f, u_) * dx

solve(a == Wext, v, bc)


# We then output the displacement and rotation fields to a XDMF file.

# In[18]:


u = v.sub(0, True)
u.rename("Displacement", "u")
theta = v.sub(1, True)
theta.rename("Rotation", "theta")
ffile.write(u, 0)
ffile.write(theta, 0)
ffile.close()


# which yields the following deformed configuration (2000 times amplification):
# 
# <center><img src="deformed_configuration.png" width="800"></center>
# 
# 
# ## References
# 
# [CAM03] Campello, E. M. B., Pimenta, P. M., & Wriggers, P. (2003). A triangular finite shell element based on a fully nonlinear shell formulation. Computational Mechanics, 31(6), 505-518.

# In[ ]: