import numpy as nm

from sfepy.base.base import assert_
from sfepy.linalg import dot_sequences
from sfepy.mechanics.tensors import dim2sym, transform_data
import sfepy.mechanics.membranes as membranes
from sfepy.terms.terms import Term

def eval_membrane_mooney_rivlin(a1, a2, mtx_c, c33, mode):
    """
    Evaluate stress or tangent stiffness of the Mooney-Rivlin membrane.

    [1] Baoguo Wu, Xingwen Du and Huifeng Tan: A three-dimensional FE
    nonlinear analysis of membranes, Computers & Structures 59 (1996),
    no. 4, 601--605.
    """
    a12 = 2.0 * a1[..., 0, 0]
    a22 = 2.0 * a2[..., 0, 0]

    sh = mtx_c.shape
    sym = dim2sym(sh[2])

    c11 = mtx_c[..., 0, 0]
    c12 = mtx_c[..., 0, 1]
    c22 = mtx_c[..., 1, 1]
    pressure = c33 * (a12 + a22 * (c11 + c22))

    if mode == 0:
        out = nm.empty((sh[0], sh[1], sym, 1))

        # S_11, S_22, S_12.
        out[..., 0, 0] = -pressure * c22 * c33 + a12 + a22 * (c22 + c33)
        out[..., 1, 0] = -pressure * c11 * c33 + a12 + a22 * (c11 + c33)
        out[..., 2, 0] = +pressure * c12 * c33 - a22 * c12

    else:
        out = nm.empty((sh[0], sh[1], sym, sym))

        dp11 = a22 * c33 - pressure * c22 * c33
        dp22 = a22 * c33 - pressure * c11 * c33
        dp12 = 2.0 * pressure * c12 * c33

        # D_11, D_22, D_33
        out[..., 0, 0] = - 2.0 * ((a22 - pressure * c22) * c22 * c33**2
                                  + c33 * c22 * dp11)
        out[..., 1, 1] = - 2.0 * ((a22 - pressure * c11) * c11 * c33**2
                                  + c33 * c11 * dp22)
        out[..., 2, 2] = - a22 + pressure * (c33 + 2.0 * c12**2 * c33**2) \
                         + c12 * c33 * dp12

        # D_21, D_31, D_32
        out[..., 1, 0] = 2.0 * ((a22 - pressure * c33
                                 - (a22 - pressure * c11) * c22 * c33**2)
                                - c33 * c11 * dp11)
        out[..., 2, 0] = 2.0 * (-pressure * c12 * c22 * c33**2
                                + c12 * c33 * dp11)
        out[..., 2, 1] = 2.0 * (-pressure * c12 * c11 * c33**2
                                + c12 * c33 * dp22)

        out[..., 0, 1] = out[..., 1, 0]
        out[..., 0, 2] = out[..., 2, 0]
        out[..., 1, 2] = out[..., 2, 1]

        # D_12, D_13, D_23
        ## out[..., 0, 1] = 2.0 * ((a22 - pressure * c33
        ##                          - (a22 - pressure * c22) * c11 * c33**2)
        ##                         - c33 * c22 * dp22)
        ## out[..., 0, 2] = 2.0 * (a22 - pressure * c22) * c12 * c33**2 \
        ##                  - c33 * c22 * dp12
        ## out[..., 1, 2] = 2.0 * (a22 - pressure * c11) * c12 * c33**2 \
        ##                  - c33 * c11 * dp12

    return out

class TLMembraneTerm(Term):
    r"""
    Mooney-Rivlin membrane with plain stress assumption.

    The membrane has a uniform initial thickness :math:`h_0` and obeys a
    hyperelastic material law with strain energy by Mooney-Rivlin: :math:`\Psi
    = a_1 (I_1 - 3) + a_2 (I_2 - 3)`.

    :Arguments:
        - material_a1 : :math:`a_1`
        - material_a2 : :math:`a_2`
        - material_h0 : :math:`h_0`
        - virtual     : :math:`\ul{v}`
        - state       : :math:`\ul{u}`
    """
    name = 'dw_tl_membrane'
    arg_types = ('material_a1', 'material_a2', 'material_h0',
                 'virtual', 'state')
    arg_shapes = {'material_a1' : '1, 1', 'material_a2' : '1, 1',
                  'material_h0' : '1, 1',
                  'virtual' : ('D', 'state'), 'state' : 'D'}
    geometries = ['3_4', '3_8']
    integration = 'surface'

    @staticmethod
    def function(out, fun, *args):
        """
        Notes
        -----
        `fun` is either `weak_function` or `eval_function` according to
        evaluation mode.
        """
        return fun(out, *args)

    @staticmethod
    def weak_function(out, a1, a2, h0, mtx_c, c33, mtx_b, mtx_t, bfg, geo,
                      fmode):
        crt = eval_membrane_mooney_rivlin(a1, a2, mtx_c, c33, fmode)

        if fmode == 0:
            bts = dot_sequences(mtx_b, crt, 'ATB')

            status = geo.integrate(out, bts * h0)
            membranes.transform_asm_vectors(out, mtx_t)

        else:
            btd = dot_sequences(mtx_b, crt, 'ATB')
            btdb = dot_sequences(btd, mtx_b)

            stress = eval_membrane_mooney_rivlin(a1, a2, mtx_c, c33, 0)

            kts =  membranes.get_tangent_stress_matrix(stress, bfg)

            mtx_k = kts + btdb

            status = geo.integrate(out, mtx_k * h0)
            membranes.transform_asm_matrices(out, mtx_t)

        return status

    @staticmethod
    def eval_function(out, a1, a2, h0, mtx_c, c33, mtx_b, mtx_t, geo,
                      term_mode, fmode):

        if term_mode == 'strain':
            out_qp = membranes.get_green_strain_sym3d(mtx_c, c33)

        elif term_mode == 'stress':
            n_el, n_qp, dm, _ = mtx_c.shape
            dim = dm + 1
            sym = dim2sym(dim)
            out_qp = nm.zeros((n_el, n_qp, sym, 1), dtype=mtx_c.dtype)

            stress = eval_membrane_mooney_rivlin(a1, a2, mtx_c, c33, 0)
            out_qp[..., 0:2, 0] = stress[..., 0:2, 0]
            out_qp[..., 3, 0] = stress[..., 2, 0]

        status = geo.integrate(out, out_qp, fmode)
        out[:, 0, :, 0] = transform_data(out.squeeze(), mtx=mtx_t)

        return status

    def __init__(self, *args, **kwargs):
        Term.__init__(self, *args, **kwargs)

        self.mtx_t = None
        self.membrane_geo = None
        self.bfg = None

    def get_fargs(self, a1, a2, h0, virtual, state,
                  mode=None, term_mode=None, diff_var=None, **kwargs):
        vv, vu = virtual, state

        sg, _ = self.get_mapping(vv)
        sd = vv.field.surface_data[self.region.name]

        if self.mtx_t is None:
            aux = membranes.describe_geometry(vu.field,
                                              self.region, self.integral)
            self.mtx_t, self.membrane_geo = aux
            # Transformed base function gradient w.r.t. material coordinates
            # in quadrature points.
            self.bfg = self.membrane_geo.bfg

        mtx_t = self.mtx_t
        bfg = self.bfg
        geo = self.membrane_geo

        # Displacements of element nodes.
        vec_u = vu.get_state_in_region(self.region)
        el_u = vec_u[sd.leconn]

        # Transform displacements to the local coordinate system.
        # u_new = T^T u
        el_u_loc = dot_sequences(el_u, mtx_t, 'AB')
        ## print el_u_loc

        mtx_c, c33, mtx_b = membranes.describe_deformation(el_u_loc, bfg)

        if mode == 'weak':
            fmode = diff_var is not None

            return (self.weak_function,
                    a1, a2, h0, mtx_c, c33, mtx_b, mtx_t, bfg, geo, fmode)

        else:
            fmode = {'eval' : 0, 'el_avg' : 1, 'qp' : 2}.get(mode, 1)

            assert_(term_mode in ['strain', 'stress'])

            return (self.eval_function,
                    a1, a2, h0, mtx_c, c33, mtx_b, mtx_t, geo, term_mode, fmode)

    def get_eval_shape(self, a1, a2, h0, virtual, state,
                       mode=None, term_mode=None, diff_var=None, **kwargs):
        n_el, n_qp, dim, n_en, n_c = self.get_data_shape(state)
        sym = dim2sym(dim)

        return (n_el, 1, sym, 1), state.dtype
