#
##
##  This file is part of pyFormex 1.0.5  (Sat Feb 16 10:40:32 CET 2019)
##  pyFormex is a tool for generating, manipulating and transforming 3D
##  geometrical models by sequences of mathematical operations.
##  Home page: http://pyformex.org
##  Project page:  http://savannah.nongnu.org/projects/pyformex/
##  Copyright 2004-2018 (C) Benedict Verhegghe (benedict.verhegghe@ugent.be)
##  Distributed under the GNU General Public License version 3 or later.
##
##  This program is free software: you can redistribute it and/or modify
##  it under the terms of the GNU General Public License as published by
##  the Free Software Foundation, either version 3 of the License, or
##  (at your option) any later version.
##
##  This program is distributed in the hope that it will be useful,
##  but WITHOUT ANY WARRANTY; without even the implied warranty of
##  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
##  GNU General Public License for more details.
##
##  You should have received a copy of the GNU General Public License
##  along with this program.  If not, see http://www.gnu.org/licenses/.
##
"""Compute inertia related quantities of geometrical models.

Inertia related quantities of a geometrical model comprise: the total mass,
the center of mass, the inertia tensor, the principal axes of inertia.

This module defines some classes to store the inertia data:

- :class:`Tensor`: a general second order tensor
- :class:`Inertia`: a specialized second order tensor for inertia data

This module also provides the basic functions to compute the inertia data
of collections of simple geometric data: points, lines, triangles, tetrahedrons.

The prefered way to compute inertia data of a geometric model is through the
:meth:`Geometry.inertia` methods.
"""
from __future__ import absolute_import, division, print_function

from pyformex import arraytools as at
from pyformex.coords import Coords
from pyformex.formex import Formex
from pyformex.coordsys import CoordSys

import numpy as np

class Tensor(np.ndarray):
    """A second order symmetric(!) tensor in 3D vector space.

    This is a new class under design. Only use for development!

    The Tensor class provides conversion between full matrix (3,3) shape
    and contracted vector (6,) shape.
    It can e.g. be used to store an inertia tensor or a stress or strain
    tensor.
    It provides methods to transform the tensor to other (cartesian)
    axes.

    Parameters:

    - `data`: array-like (float) of shape (3,3) or (6,)
    - `symmetric`: bool. If True (default), the tensor is forced to be
      symmetric by averaging the off-diagonal elements.
    - `cs`: CoordSys. The coordinate system of the tensor.

    Properties: a Tensor T has the following properties:

    - T.xx, T.xy, T.xz, T.yx, T.yy, T.yz, T.zx, T.zy, T.zz: aliases for
      the nine components of the tensor
    - T.contracted: the (6,) shaped contracted array with independent
      values of the tensor
    - T.tensor: the full tensor as an (3,3) array

    Discussion:

    - inertia and stres/strain tensors transform in the same way
      on rotations of axes. But differently on translations! Should we
      therefore store the purpose of the tensor??

      * Propose to leave it to the user to know what he is doing.
      * Propose to have a separate class Inertia derived from Tensor,
        which implements computing the inertia tensor and translation.

    - should we allow non-symmetrical tensors? Then what with principal?

      * Propose to silently allow non-symm. Result of functions is what it is.
        Again, suppose the user knows what he is doing.

    Example:

        >>> t = Tensor([1,2,3,4,5,6])
        >>> print(t)
        [[ 1. 6. 5.]
         [ 6. 2. 4.]
         [ 5. 4. 3.]]
        >>> print(t.contracted)
        [ 1. 2. 3. 4. 5. 6.]
        >>> s = Tensor(t)
        >>> print(s)
        [[ 1. 6. 5.]
         [ 6. 2. 4.]
         [ 5. 4. 3.]]

    """

    _contracted_order = [(0, 1, 2, 1, 2, 0), (0, 1, 2, 2, 0, 1)]
    _contracted_index = np.array([
        [ 0, 5, 4],
        [ 5, 1, 3],
        [ 4, 3, 2],
        ])


    def __new__(clas,data=None,symmetric=True,cs=None):
        """Create a new Tensor instance"""
        try:
            data = at.checkArray(data,shape=(3,3),kind='f',allow='if')
        except:
            try:
                data = at.checkArray(data,shape=(6,),kind='f',allow='if')
            except:
                raise ValueError("Data should have shape (3,3) or (6,)")
            data = data[Tensor._contracted_index]
        ar = data.view(clas)
        if cs is None:
            cs = CoordSys()
        if not(isinstance(cs,CoordSys)):
            raise ValueError('Wrong Coordinate System')
        ar.cs = cs
        return ar


    def __array_finalize__(self, obj):
        """Finalize the new Matrix object.

        When a class is derived from numpy.ndarray and the constructor (the
        :meth:`__new__` method) defines new attributes, these atttributes
        need to be reset in this method.
        """
        if obj is None: return
        self.cs = getattr(obj, 'cs', CoordSys())
        #~ pass  # currently no new attributes defined in __new__
        # self._newattr = getattr(obj, '_newattr', None)

    @property
    def xx(self):
        return self[0,0]
    @property
    def yy(self):
        return self[1,1]
    @property
    def zz(self):
        return self[2,2]
    @property
    def yz(self):
        return self[1,2]
    @property
    def zx(self):
        return self[0,2]
    @property
    def xy(self):
        return self[0,1]

    zy = yz
    xz = zx
    yx = xy


    @property
    def contracted(self):
        """Returned the symmetric tensor data as a numpy array with shape (6,)

        """
        return self.sym[Tensor._contracted_order]

    @property
    def tensor(self):
        """Returned the tensor data as a numpy array with shape (3,3)"""
        return np.asarray(self)

    @property
    def sym(self):
        """Return the symmetric part of the tensor."""
        return (self+self.T) / 2

    @property
    def asym(self):
        """Return the antisymmetric part of the tensor."""
        return (self-self.T) / 2


    def principal(self,sort=True,right_handed=True):
        """Returns the principal values and axes of the inertia tensor.

        Parameters:

        - `sort`: bool. If True (default), the return values are sorted
          in order of decreasing principal values. Otherwise they are
          unsorted.
        - `right_handed`: bool. If True (default), the returned axis vectors
          are guaranteed to form a right-handed coordinate system.
          Otherwise, lef-handed systems may result)

        Returns a tuple (prin,axes) where

        - `prin`: is a (3,) array with the principal values,
        - `axes`: is a (3,3) array with the rotation matrix that
          rotates the global axes to the principal axes. This also
          means that the rows of axes are the unit vectors along
          the principal directions.

        Example:

        >>> t = Tensor([-19., 4.6, -8.3, 11.8, 6.45, -4.7 ])
        >>> p,a = t.principal()
        >>> print(p)
        [ 11.62  -9.   -25.32]
        >>> print(a)
        [[-0.03  0.86  0.5 ]
         [-0.62  0.38 -0.69]
         [-0.78 -0.33  0.53]]
        """
        #
        # Req: requires numpy >= 1.8.0
        #
        prin, axes = np.linalg.eig(self.tensor)
        axes = axes.transpose()  # put the eigenvectors rowwise.
        if sort:
            s = prin.argsort()[::-1]
            prin, axes = prin[s], axes[s]
        if right_handed and not at.allclose(at.normalize(np.cross(axes[0], axes[1])), axes[2]):
            axes[2] = -axes[2]
        return prin, axes


    def rotate(self,rot):
        """Transform the tensor on coordinate system rotation.

        Note: for an inertia tensor, the inertia should have been
        computed around axes through the center of mass.
        See also translate.

        Example:

        >>> t = Tensor([-19., 4.6, -8.3, 11.8, 6.45, -4.7 ])
        >>> p,a = t.principal()
        >>> print(t.rotate(np.linalg.linalg.inv(a)))
        [[ 11.62   0.     0.  ]
         [ -0.    -9.    -0.  ]
         [  0.    -0.   -25.32]]

        """
        rot = at.checkArray(rot,shape=(3,3),kind='f')
        return Tensor(at.atba(rot,self),self.cs.rotate(rot))


class Inertia(Tensor):
    """A class for storing the inertia tensor of an array of points.

    Parameters:

    - `X`: a Coords with shape (npoints,3). Shapes (...,3) are accepted
      but will be reshaped to (npoints,3).
    - `mass`: optional, (npoints,) float array with the mass of the points.
      If omitted, all points have mass 1.

      The result is a tuple of two float arrays:

      - the center of gravity: shape (3,)
      - the inertia tensor: shape (6,) with the following values (in order):
        Ixx, Iyy, Izz, Iyz, Izx, Ixy

    Example:

    >>> from .elements import Tet4
    >>> X = Tet4.vertices
    >>> print(X)
    [[ 0. 0. 0.]
     [ 1. 0. 0.]
     [ 0. 1. 0.]
     [ 0. 0. 1.]]
    >>> I = X.inertia()
    >>> print(I)
    [[ 1.5  0.25 0.25]
     [ 0.25 1.5  0.25]
     [ 0.25 0.25 1.5 ]]
    >>> print(I.ctr)
    [ 0.25 0.25 0.25]
    >>> print(I.mass)
    4.0
    >>> print(I.translate(-I.ctr))
    [[ 2. 0. 0.]
     [ 0. 2. 0.]
     [ 0. 0. 2.]]

    """

    def __new__(clas,data,mass,ctr):
        """Create a new Tensor instance"""
        ar = Tensor.__new__(clas,data)
        # We need mass and ctr!
        ar.mass = float(mass)
        ar.ctr = Coords(at.checkArray(ctr,shape=(3,),kind='f'))
        return ar


    def translate(self,trl,toG=False):
        """Return the inertia tensor around axes translated over vector trl.

        Parameters:

        - `trl`: arraylike (3,). Distance vector from the center of mass
          to the new reference point.
        - `toG`: bool. If False (default) the inertia tensor is translated to the
          the  new reference point, otherwise it will be translated to its center
          of mass

        """
        trl = at.checkArray(trl,shape=(3,),kind='f')
        trf = -np.dot(trl.reshape(3,-1),trl.reshape(-1,3))
        ind = np.diag_indices(3)
        trf[ind] += np.dot(trl,trl.T)
        if toG:
            trf = -trf
        return self.tensor + self.mass * trf


    def translateTo(self,ref,toG=False):
        """Return the inertia tensor around axes translated to the reference
        point ref.

        Parameters:

        - `ref`: arraylike (3,). The new reference point coordinates.
        - `toG`: bool. If False (default) the inertia tensor is translated to the
          the  new reference point, otherwise it will be translated to its center
          of mass

        """

        trl = at.checkArray(ref,shape=(3,),kind='f') - self.ctr
        return self.translate(self,trl,toG=toG)


    def toCS(self,cs):
        """Transform the coordinates to another CoordSys.

        """
        raise ImplementationError("Inertia.toCS is currently not implemented")
        # We need to check the following (after transposing CoordSys.rot)
        return self.tensor.translateTo(cs.trl) + self.tensor.rotate(np.dot(self.CS.rot.T,cs.rot))


def point_inertia(X,mass=None,center_only=False):
    """Compute the total mass, center of mass and inertia tensor mass points.

    Parameters:

    - `X`: a Coords with shape (npoints,3). Shapes (...,3) are accepted
      but will be reshaped to (npoints,3).
    - `mass`: optional, (npoints,) float array with the mass of the points.
      If omitted, all points have mass 1.
    - `center_only`: bool: if True, only returns the total mass and center
      of mass.

    Returns a tuple (M,C,I) where M is the total mass of all points, C is
    the center of mass, and I is the inertia tensor in the central coordinate
    system, i.e. a coordinate system with axes paralle to the global axes
    but origin at the (computed) center of mass. If `center_only` is True,
    returns the tuple (M,C) only. On large models this is more effective
    in case you do not need the inertia tensor.

    """
    X = Coords(X).reshape(-1, 3)
    npoints = X.shape[0]
    if mass is not None:
        mass = at.checkArray(mass,shape=(npoints,),kind='f')
        mass = mass.reshape(-1,1)
        M = mass.sum()
        C = (X*mass).sum(axis=0) / M
    else:
        M = float(npoints)
        C = X.mean(axis=0)
    if center_only:
        return M,C

    Xc = X - C
    x, y, z = Xc[:, 0], Xc[:, 1], Xc[:, 2]
    xx, yy, zz, yz, zx, xy = x*x, y*y, z*z, y*z, z*x, x*y
    I = np.column_stack([ yy+zz, zz+xx, xx+yy, -yz, -zx, -xy ])
    if mass is not None:
        I *= mass
    I = I.sum(axis=0)
    return M,C,I


def surface_volume(x,pt=None):
    """Return the volume inside a 3-plex Formex.

    - `x`: an (ntri,3,3) shaped float array, representing ntri triangles.
    - `pt`: a point in space. If unspecified, it is taken equal to the
      origin of the global coordinate system ([0.,0.,0.]).

    Returns an (ntri) shaped array with the volume of the tetrahedrons formed
    by the triangles and the point `pt`. Triangles with an outer normal
    pointing away from `pt` will generate positive tetrahral volumes, while
    triangles having `pt` at the side of their positive normal will generate
    negative volumes. In any case, if `x` represents a closed surface,
    the algebraic sum of all the volumes is the total volume inside the surface.
    """
    x = at.checkArray(x,shape=(-1,3,3),kind='f')
    if pt is not None:
        x -= pt
    a, b, c = [ x[:,i,:] for i in range(3) ]
    d = np.cross(b, c)
    e = (a*d).sum(axis=-1)
    return e / 6


def surface_volume_inertia(x,center_only=False):
    """Return the inertia of the volume inside a 3-plex Formex.

    - `x`: an (ntri,3,3) shaped float array, representing ntri triangles.

    This uses the same algorithm as tetrahedral_inertia using [0.,0.,0.]
    as the 4-th point for each tetrahedron.

    Returns a tuple (V,C,I) where V is the total volume,
    C is the center of mass (3,) and I is the inertia tensor (6,) of the
    tetrahedral model.

    Example:

    >>> from .simple import sphere
    >>> S = sphere(4).toFormex()
    >>> V,C,I = surface_volume_inertia(S.coords)
    >>> print(V,C,I)
    4.04701 [-0. -0. -0.] [ 1.58  1.58  1.58 -0.    0.    0.  ]

    """
    def K(x,y):
        x1,x2,x3 = x[:,0], x[:,1], x[:,2]
        y1,y2,y3 = y[:,0], y[:,1], y[:,2]
        return x1 * (y1+y2+y3) + \
               x2 * (   y2+y3) + \
               x3 * (      y3)

    x = at.checkArray(x,shape=(-1,3,3),kind='f')
    v = surface_volume(x)
    V = v.sum()
    c = x.sum(axis=1) / 4.  # 4-th point is 0.,0.,0.
    C = (c*v[:,np.newaxis]).sum(axis=0) / V
    if center_only:
        return V,C

    x -= C
    aa = 2 * K(x,x) * v.reshape(-1,1)
    aa = aa.sum(axis=0)
    a0 = aa[1] + aa[2]
    a1 = aa[0] + aa[2]
    a2 = aa[0] + aa[1]
    x0,x1,x2 = x[...,0],x[...,1],x[...,2]
    a3 = (( K(x1,x2) + K(x2,x1) ) * v).sum(axis=0)
    a4 = (( K(x2,x0) + K(x0,x2) ) * v).sum(axis=0)
    a5 = (( K(x0,x1) + K(x1,x0) ) * v).sum(axis=0)
    I = np.array([a0,a1,a2,a3,a4,a5]) / 20.
    return V,C,I


def tetrahedral_volume(x):
    """Compute the volume of tetrahedrons.

    - `x`: an (ntet,4,3) shaped float array, representing ntet tetrahedrons.

    Returns an (ntet,) shaped array with the volume of the tetrahedrons.
    Depending on the ordering of the points, this volume may be positive
    or negative. It will be positive if point 4 is on the side of the positive
    normal formed by the first 3 points.
    """
    x = at.checkArray(x,shape=(-1,4,3),kind='f')
    a, b, c = [ x[:,i,:] - x[:,3,:] for i in range(3) ]
    d = np.cross(b, c)
    e = (a*d).sum(axis=-1)
    return -e / 6


def tetrahedral_inertia(x,density=None,center_only=False):
    """Return the inertia of the volume of a 4-plex Formex.

    Parameters:

    - `x`: an (ntet,4,3) shaped float array, representing ntet tetrahedrons.
    - `density`: optional mass density (ntet,) per tetrahedron
    - `center_only`: bool. If True, returns only the total volume, total mass
      and center of gravity. This may be used on large models when only these
      quantities are required.

    Returns a tuple (V,M,C,I) where V is the total volume, M is the total mass,
    C is the center of mass (3,) and I is the inertia tensor (6,) of the
    tetrahedral model.

    Formulas for inertia were based on F. Tonon, J. Math & Stat, 1(1):8-11,2005

    Example:

    >>> x = Coords([
    ...     [  8.33220, -11.86875,  0.93355 ],
    ...     [  0.75523,   5.00000, 16.37072 ],
    ...     [ 52.61236,   5.00000, -5.38580 ],
    ...     [  2.000000,  5.00000,  3.00000 ],
    ...     ])
    >>> F = Formex([x])
    >>> print(tetrahedral_center(F.coords))
    [ 15.92  0.78  3.73]
    >>> print(tetrahedral_volume(F.coords))
    [ 1873.23]
    >>> print(*tetrahedral_inertia(F.coords))
    1873.23 1873.23 [ 15.92  0.78  3.73] [ 43520.32 194711.28 191168.77   4417.66 -46343.16  11996.2 ]

    """
    def K(x,y):
        x1,x2,x3,x4 = x[:,0], x[:,1], x[:,2], x[:,3]
        y1,y2,y3,y4 = y[:,0], y[:,1], y[:,2], y[:,3]
        return x1 * (y1+y2+y3+y4) + \
               x2 * (   y2+y3+y4) + \
               x3 * (      y3+y4) + \
               x4 * (         y4)

    x = at.checkArray(x,shape=(-1,4,3),kind='f')
    v = tetrahedral_volume(x)
    V = v.sum()
    if density:
        v *= density
    c = Formex(x).centroids()
    M = v.sum()
    C = (c*v[:,np.newaxis]).sum(axis=0) / M
    if center_only:
        return V,M,C

    x -= C
    aa = 2 * K(x,x) * v.reshape(-1,1)
    aa = aa.sum(axis=0)
    a0 = aa[1] + aa[2]
    a1 = aa[0] + aa[2]
    a2 = aa[0] + aa[1]
    x0,x1,x2 = x[...,0],x[...,1],x[...,2]
    a3 = (( K(x1,x2) + K(x2,x1) ) * v).sum(axis=0)
    a4 = (( K(x2,x0) + K(x0,x2) ) * v).sum(axis=0)
    a5 = (( K(x0,x1) + K(x1,x0) ) * v).sum(axis=0)
    I = np.array([a0,a1,a2,a3,a4,a5]) / 20.
    return V,M,C,I


def tetrahedral_center(x,density=None):
    """Compute the center of mass of a collection of tetrahedrons.

    - `x`: an (ntet,4,3) shaped float array, representing ntet tetrahedrons.
    - `density`: optional mass density (ntet,) per tetrahedron. Default 1.

    Returns a (3,) shaped array with the center of mass.
    """
    return tetrahedral_inertia(x,density=None,center_only=True)[2]


# Kept for compatibility
inertia = point_inertia


# End
