# Source code for plugins.isopar

    # $Id$
    ##
    ##  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/.
    ##

    """Isoparametric transformations

    """
    from __future__ import absolute_import, division, print_function

    from pyformex import arraytools as at
    from pyformex.coords import Coords
    from .polynomial import Polynomial

    import numpy as np

    #
    # This should probably be moved to polynomial
    #
    def evaluate(atoms,x,y=0,z=0):
        """Build a matrix of functions of coords.

        - `atoms`: a list of text strings representing a mathematical function of
          `x`, and possibly of `y` and `z`.
        - `x`, `y`, `z`: a list of x- (and optionally y-, z-) values at which the
          `atoms` will be evaluated. The lists should have the same length.

        Returns a matrix with `nvalues` rows and `natoms` colums.
        """
        aa = np.zeros((len(x), len(atoms)), at.Float)
        for k, a in enumerate(atoms):
            aa[:, k] = eval(a)
        return aa

    def exponents(n,layout='lag'):
        """Create tuples of polynomial exponents.

        This function creates the exponents of polynomials in 1 to 3 dimensions
        which can be used to construct interpolation function over lagrangian,
        triangular or serendipity grids.

        Parameters:

        - `n`: a tuple of 1 to 3 integers, specifying the degree of the polynomials
          in the x up to z directions. For a lagrangian layout, this is one less
          than the number of points in each direction.
        - `layout`: string, specifying the layout of grid and the selection of
          monomials to be used. Should be one of 'lagrangian', 'triangular',
          'serendipity' or 'border'. The string can be abbreviated to its
          first 3 characters.

        Returns an integer array of shape (ndim,npoints), where ndim = len(n)
        and npoints depends on the layout:

        - lagrangian: npoints = prod(n). The point layout is a rectangular
          lagrangian grid form by n[i] points in direction i. As an example,
          specifying n=(3,2) uses a grid of 3 points in x-direction and 2 points
          in y-direction.
        - triangular: requires that all values in n are equal. For ndim=2, the
          number of points is n*(n+1)//2.
        - border: this is like the lagrangian grid with all internal points
          removed. For ndim=2, we have npoints = 2 * sum(n) - 4. For ndim=3 we
          have npoints = 2 * sum(nx*ny+ny*nz+nz*nx) - 4 * sum(n) + 8.
          Thus n=(3,3,3) will yield 2*3*3*3 - 4*(3+3+3) + 8 = 26
        - serendipity: tries to use only the corner and edge nodes, but uses
          a convex domain of the monomials. This may require some nodes inside
          the faces or the volume. Currently works up to (4,4) in 2D or (3,3,3)
          in 3D.

        """
        n = at.checkArray(n, (-1,), 'i')
        ndim = n.shape[0]
        if ndim < 1 or ndim > 3:
            raise RuntimeError("Expected a 1..3 length tuple")

        layout = layout[:3]
        if layout in ['tri', 'ser']:
            if not (n == n[0]).all():
                raise RuntimeError("For triangular and serendipity grids, all axes should have the same number of points")

        # First create the full lagrangian set
        exp = np.indices(n+1).reshape(ndim, -1).transpose()

        if layout != 'lag':
            if layout == 'tri':
                # Select by total maximal degree
                ok = exp.sum(axis=-1) <= n[0]
            elif layout == 'bor':
                # Select if any value is not higher than 1
                ok = (exp <= 1).any(axis=-1)
            elif layout == 'ser':
                if len(n) > 2:
                    #print(exp.sum(axis=-1))
                    #print(n[0] + len(n) - 1)
                    npts = 8 + 12*(n[0] - 1) # minimal number of points
                    sdeg = n[0] + 1
                    ok = exp.sum(axis=-1) <= sdeg
                    if ok.sum() < npts:
                        #print("Need at least %s" % npts)
                        sdeg += 1
                        ok = exp.sum(axis=-1) <= sdeg
                        ok1 = (exp == n[0]).sum(axis=-1) <= 1
                        ok = ok*ok1
                else:
                    npts = 4 + 4*(n[0] - 1) # minimal number of points
                    ok = exp.sum(axis=-1) <= n[0] + 1
                if ok.sum() < npts:
                    raise ValueError("No solution for eltype %s %s" % (layout, n))
            else:
                raise RuntimeError("Unknown layout %s" % layout)
            exp = exp[ok]
        return exp

    def interpoly(n,layout='lag'):
        """Create an interpolation polynomial

        parameters are like for exponents.

        Returns a Polynomial that can be used for interpolation over
        the element.
        """
        exp = exponents(n, layout)
        return Polynomial(exp)

    class Isopar(object):
        """A class representing an isoparametric transformation

            eltype is one of the keys in Isopar.isodata
            coords and oldcoords can be either arrays, Coords or Formex instances,
            but should be of equal shape, and match the number of atoms in the
            specified transformation type

        The following three formulations are equivalent ::

           trf = Isopar(eltype,coords,oldcoords)
           G = F.isopar(trf)

           trf = Isopar(eltype,coords,oldcoords)
           G = trf.transform(F)

           G = isopar(F,eltype,coords,oldcoords)

        """
        # Store the interpolation polynomials corresponding with
        # some element
        isodata = {
            }

        isodata_alias = {
            'line2': 'lag-1',
            'line3': 'lag-2',
            'line4': 'lag-3',
            'tri3': 'tri-1-1',
            'tri6': 'tri-2-2',
            'tri10': 'tri-3-3',
            'quad4': 'lag-1-1',
            'quad9': 'lag-2-2',
            'quad16': 'lag-3-3',
            'quad8': 'bor-2-2',
            'quad12': 'bor-3-3',
            'quad13': 'ser-3-3',
            'tet4': 'tri-1-1-1',
            'tet10': 'tri-2-2-2',
            'hex8': 'lag-1-1-1',
            'hex20': 'ser-2-2-2',
            'hex27': 'lag-2-2-2',
            'hex36': 'lag-2-2-3',
            'hex64': 'lag-3-3-3',
            }

        def __init__(self, eltype, coords, oldcoords):
            """Create an isoparametric transformation.

            `eltype`: string: either one of the keys in the isodata dictionary,
              or a string in the following format: layout-nx-ny-nz

            """
            if eltype not in Isopar.isodata:
                if eltype in Isopar.isodata_alias:
                    eltype = Isopar.isodata_alias[eltype]
            if eltype not in Isopar.isodata:
                s = eltype.split('-')
                if s[0] in [ 'lag', 'tri', 'bor', 'ser' ]:
                    n = [int(v) for v in s[1:]]
                    #
                    # It might be better to just store (ndim,atoms)
                    # if we use string atoms to evaluate
                    #
                    Isopar.isodata[eltype] = interpoly(n, s[0])
                else:
                    raise RuntimeError("Unknown eltype %s")

            poly = Isopar.isodata[eltype]
            coords = coords.view().reshape(-1, 3)
            oldcoords = oldcoords.view().reshape(-1, 3)
            aa = poly.evalAtoms(oldcoords[:, :poly.ndim])
            ab = np.linalg.solve(aa, coords)
            self.eltype = eltype
            self.trf = ab

        def transform(self, X):
            """Apply isoparametric transform to a set of coordinates.

            Returns a Coords array with same shape as X
            """
            if not isinstance(X, Coords):
                try:
                    X = Coords(X)
                except:
                    raise ValueError("Expected a Coords object as argument")

            poly = Isopar.isodata[self.eltype]
            ndim = poly.ndim
            aa = poly.evalAtoms(X.reshape(-1, 3)[:, :ndim])
            xx = np.dot(aa, self.trf).reshape(X.shape)
            if ndim < 3:
                xx[..., ndim:] += X[..., ndim:]
            return X.__class__(xx)

    # End