r"""
Piezo-elasticity problem - linear elastic material with piezoelectric
effects.

Find :math:`\ul{u}`, :math:`\phi` such that:

.. math::
    - \omega^2 \int_{Y} \rho\ \ul{v} \cdot \ul{u}
    + \int_{Y} D_{ijkl}\ e_{ij}(\ul{v}) e_{kl}(\ul{u})
    - \int_{Y_2} g_{kij}\ e_{ij}(\ul{v}) \nabla_k \phi
    = 0
    \;, \quad \forall \ul{v} \;,

    \int_{Y_2} g_{kij}\ e_{ij}(\ul{u}) \nabla_k \psi
    + \int_{Y} K_{ij} \nabla_i \psi \nabla_j \phi
    = 0
    \;, \quad \forall \psi \;,

where

.. math::
    D_{ijkl} = \mu (\delta_{ik} \delta_{jl}+\delta_{il} \delta_{jk}) +
    \lambda \ \delta_{ij} \delta_{kl}
    \;.
"""
from __future__ import absolute_import
import os
import numpy as nm

from sfepy import data_dir
from sfepy.discrete.fem import MeshIO
from sfepy.mechanics.matcoefs import stiffness_from_lame
import six

def post_process(out, pb, state, extend=False):
    """
    Calculate and output the strain and stresses for the given state.
    """
    from sfepy.base.base import Struct
    from sfepy.discrete.fem import extend_cell_data

    ev = pb.evaluate
    strain = ev('ev_cauchy_strain.i.Y(u)', mode='el_avg')
    stress = ev('ev_cauchy_stress.i.Y(inclusion.D, u)', mode='el_avg')

    piezo = -ev('ev_piezo_stress.i.Y2(inclusion.coupling, phi)',
                mode='el_avg')
    piezo = extend_cell_data(piezo, pb.domain, 'Y2', val=0.0)

    piezo_strain = ev('ev_piezo_strain.i.Y(inclusion.coupling, u)',
                      mode='el_avg')

    out['cauchy_strain'] = Struct(name='output_data', mode='cell',
                                  data=strain, dofs=None)
    out['elastic_stress'] = Struct(name='output_data', mode='cell',
                                   data=stress, dofs=None)
    out['piezo_stress'] = Struct(name='output_data', mode='cell',
                                 data=piezo, dofs=None)
    out['piezo_strain'] = Struct(name='output_data', mode='cell',
                                 data=piezo_strain, dofs=None)
    out['total_stress'] = Struct(name='output_data', mode='cell',
                                 data=stress + piezo, dofs=None)

    return out

filename_mesh = data_dir + '/meshes/2d/special/circle_in_square.mesh'
## filename_mesh = data_dir + '/meshes/2d/special/circle_in_square_small.mesh'
## filename_mesh = data_dir + '/meshes/3d/special/cube_sphere.mesh'
## filename_mesh = data_dir + '/meshes/2d/special/cube_cylinder.mesh'

omega = 1
omega_squared = omega**2

conf_dir = os.path.dirname(__file__)
io = MeshIO.any_from_filename(filename_mesh, prefix_dir=conf_dir)
bbox, dim = io.read_bounding_box(ret_dim=True)

geom = {3 : '3_4', 2 : '2_3'}[dim]

x_left, x_right = bbox[:,0]

options = {
    'post_process_hook' : 'post_process',
}

regions = {
    'Y' : 'all',
    'Y1' : 'cells of group 1',
    'Y2' : 'cells of group 2',
    'Y2_Surface': ('r.Y1 *v r.Y2', 'facet'),
    'Left' : ('vertices in (x < %f)' % (x_left + 1e-3), 'facet'),
    'Right' : ('vertices in (x > %f)' % (x_right - 1e-3), 'facet'),
}

fields = {
    'displacement' : ('real', dim, 'Y', 1),
    'potential' : ('real', 1, 'Y', 1),
}

variables = {
    'u' : ('unknown field', 'displacement', 0),
    'v' : ('test field', 'displacement', 'u'),
    'phi' : ('unknown field', 'potential', 1),
    'psi' : ('test field', 'potential', 'phi'),
}

ebcs = {
    'u1' : ('Left', {'u.all' : 0.0}),
    'u2' : ('Right', {'u.0' : 0.1}),
    'phi' : ('Y2_Surface', {'phi.all' : 0.0}),
}

def get_inclusion_pars(ts, coor, mode=None, **kwargs):
    """TODO: implement proper 3D -> 2D transformation of constitutive
    matrices."""
    if mode == 'qp':
        _, dim = coor.shape
        sym = (dim + 1) * dim // 2

        dielectric = nm.eye(dim, dtype=nm.float64)
        # !!!
        coupling = nm.ones((dim, sym), dtype=nm.float64)
        #    coupling[0,1] = 0.2

        out = {
            # Lame coefficients in 1e+10 Pa.
            'D' : stiffness_from_lame(dim=2, lam=0.1798, mu=0.148),
            # dielectric tensor
            'dielectric' : dielectric,
            # piezoelectric coupling
            'coupling' : coupling,
            'density' : nm.array([[0.1142]]), # in 1e4 kg/m3
        }

        for key, val in six.iteritems(out):
            out[key] = val[None, ...]

        return out

materials = {
    'inclusion' : (None, 'get_inclusion_pars')
}

functions = {
    'get_inclusion_pars' : (get_inclusion_pars,),
}

integrals = {
    'i' : 2,
}

equations = {
    '1' : """- %f * dw_volume_dot.i.Y(inclusion.density, v, u)
             + dw_lin_elastic.i.Y(inclusion.D, v, u)
             - dw_piezo_coupling.i.Y2(inclusion.coupling, v, phi)
           = 0""" % omega_squared,
    '2' : """dw_piezo_coupling.i.Y2(inclusion.coupling, u, psi)
           + dw_diffusion.i.Y(inclusion.dielectric, psi, phi)
           = 0""",
}

solvers = {
    'ls' : ('ls.scipy_direct', {}),
    'newton' : ('nls.newton',
                {'i_max'      : 1,
                 'eps_a'      : 1e-10,
    }),
}

[((74, 11, 74, 36), 'os.path.dirname', 'os.path.dirname', ({(74, 27, 74, 35): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((75, 5, 75, 65), 'sfepy.discrete.fem.MeshIO.any_from_filename', 'MeshIO.any_from_filename', (), '', False, 'from sfepy.discrete.fem import MeshIO\n'), ((48, 12, 48, 61), 'sfepy.discrete.fem.extend_cell_data', 'extend_cell_data', (), '', False, 'from sfepy.discrete.fem import extend_cell_data\n'), ((53, 27, 54, 57), 'sfepy.base.base.Struct', 'Struct', (), '', False, 'from sfepy.base.base import Struct\n'), ((55, 28, 56, 58), 'sfepy.base.base.Struct', 'Struct', (), '', False, 'from sfepy.base.base import Struct\n'), ((57, 26, 58, 55), 'sfepy.base.base.Struct', 'Struct', (), '', False, 'from sfepy.base.base import Struct\n'), ((59, 26, 60, 62), 'sfepy.base.base.Struct', 'Struct', (), '', False, 'from sfepy.base.base import Struct\n'), ((61, 26, 62, 64), 'sfepy.base.base.Struct', 'Struct', (), '', False, 'from sfepy.base.base import Struct\n'), ((120, 21, 120, 50), 'numpy.eye', 'nm.eye', (), '', True, 'import numpy as nm\n'), ((122, 19, 122, 56), 'numpy.ones', 'nm.ones', (), '', True, 'import numpy as nm\n'), ((135, 24, 135, 42), 'six.iteritems', 'six.iteritems', ({(135, 38, 135, 41): 'out'}, {}), '(out)', False, 'import six\n'), ((127, 18, 127, 66), 'sfepy.mechanics.matcoefs.stiffness_from_lame', 'stiffness_from_lame', (), '', False, 'from sfepy.mechanics.matcoefs import stiffness_from_lame\n'), ((132, 24, 132, 44), 'numpy.array', 'nm.array', ({(132, 33, 132, 43): '[[0.1142]]'}, {}), '([[0.1142]])', True, 'import numpy as nm\n')]