# -*- coding: utf-8 -*-
from dewloosh.geom.polyhedron import HexaHedron
from dewloosh.math.numint import GaussPoints as Gauss
from dewloosh.geom.utils import cells_coords
from numba import njit, prange
import numpy as np
from numpy import ndarray
__cache = True


@njit(nogil=True, cache=__cache)
def monoms_H8(pcoord: np.ndarray):
    r, s, t = pcoord
    return np.array([1, r, s, t, r*s, r*t, s*t, r*s*t])


@njit(nogil=True, cache=__cache)
def shp_H8(pcoord):
    r, s, t = pcoord
    return np.array([-0.125*r*s*t + 0.125*r*s + 0.125*r*t - 0.125*r +
                     0.125*s*t - 0.125*s - 0.125*t + 0.125,
                     0.125*r*s*t - 0.125*r*s - 0.125*r*t + 0.125*r +
                     0.125*s*t - 0.125*s - 0.125*t + 0.125,
                     -0.125*r*s*t + 0.125*r*s - 0.125*r*t + 0.125*r -
                     0.125*s*t + 0.125*s - 0.125*t + 0.125,
                     0.125*r*s*t - 0.125*r*s + 0.125*r*t - 0.125*r -
                     0.125*s*t + 0.125*s - 0.125*t + 0.125,
                     0.125*r*s*t + 0.125*r*s - 0.125*r*t - 0.125*r -
                     0.125*s*t - 0.125*s + 0.125*t + 0.125,
                     -0.125*r*s*t - 0.125*r*s + 0.125*r*t + 0.125*r -
                     0.125*s*t - 0.125*s + 0.125*t + 0.125,
                     0.125*r*s*t + 0.125*r*s + 0.125*r*t + 0.125*r +
                     0.125*s*t + 0.125*s + 0.125*t + 0.125,
                     -0.125*r*s*t - 0.125*r*s - 0.125*r*t - 0.125*r +
                     0.125*s*t + 0.125*s + 0.125*t + 0.125]
                    )


@njit(nogil=True, parallel=True, cache=__cache)
def shape_function_matrix_H8(pcoord: np.ndarray):
    eye = np.eye(3, dtype=pcoord.dtype)
    shp = shp_H8(pcoord)
    res = np.zeros((3, 24), dtype=pcoord.dtype)
    for i in prange(8):
        res[:, i*3: (i+1) * 3] = eye*shp[i]
    return res


@njit(nogil=True, cache=__cache)
def dshp_H8(pcoord):
    r, s, t = pcoord
    return np.array(
        [[-0.125*s*t + 0.125*s + 0.125*t - 0.125,
         -0.125*r*t + 0.125*r + 0.125*t - 0.125,
         -0.125*r*s + 0.125*r + 0.125*s - 0.125],
         [0.125*s*t - 0.125*s - 0.125*t + 0.125,
         0.125*r*t - 0.125*r + 0.125*t - 0.125,
         0.125*r*s - 0.125*r + 0.125*s - 0.125],
         [-0.125*s*t + 0.125*s - 0.125*t + 0.125,
         -0.125*r*t + 0.125*r - 0.125*t + 0.125,
         -0.125*r*s - 0.125*r - 0.125*s - 0.125],
         [0.125*s*t - 0.125*s + 0.125*t - 0.125,
         0.125*r*t - 0.125*r - 0.125*t + 0.125,
         0.125*r*s + 0.125*r - 0.125*s - 0.125],
         [0.125*s*t + 0.125*s - 0.125*t - 0.125,
         0.125*r*t + 0.125*r - 0.125*t - 0.125,
         0.125*r*s - 0.125*r - 0.125*s + 0.125],
         [-0.125*s*t - 0.125*s + 0.125*t + 0.125,
         -0.125*r*t - 0.125*r - 0.125*t - 0.125,
         -0.125*r*s + 0.125*r - 0.125*s + 0.125],
         [0.125*s*t + 0.125*s + 0.125*t + 0.125,
         0.125*r*t + 0.125*r + 0.125*t + 0.125,
         0.125*r*s + 0.125*r + 0.125*s + 0.125],
         [-0.125*s*t - 0.125*s - 0.125*t - 0.125,
         -0.125*r*t - 0.125*r + 0.125*t + 0.125,
         -0.125*r*s - 0.125*r + 0.125*s + 0.125]]
    )


@njit(nogil=True, parallel=True, cache=__cache)
def dshp_H8_bulk(pcoords: ndarray):
    nP = pcoords.shape[0]
    res = np.zeros((nP, 8, 3), dtype=pcoords.dtype)
    for iP in prange(nP):
        res[iP] = dshp_H8(pcoords[iP])
    return res


@njit(nogil=True, parallel=True, fastmath=True, cache=__cache)
def volumes_H8(ecoords: np.ndarray, qpos: np.ndarray,
               qweight: np.ndarray):
    nE = ecoords.shape[0]
    volumes = np.zeros(nE, dtype=ecoords.dtype)
    nQ = len(qweight)
    for iQ in range(nQ):
        dshp = dshp_H8(qpos[iQ])
        for i in prange(nE):
            jac = ecoords[i].T @ dshp
            djac = np.linalg.det(jac)
            volumes[i] += qweight[iQ] * djac
    return volumes


class H8(HexaHedron):
    """
    8-node isoparametric hexahedron.

    top        
    7--6  
    |  |
    4--5

    bottom
    3--2  
    |  |
    0--1

    """

    @classmethod
    def lcoords(cls, *args, **kwargs):
        return np.array([[-1., -1., -1],
                         [1., -1., -1.],
                         [1., 1., -1.],
                         [-1., 1., -1.],
                         [-1., -1., 1.],
                         [1., -1., 1.],
                         [1., 1., 1.],
                         [-1., 1., 1.]])

    @classmethod
    def lcenter(cls, *args, **kwargs):
        return np.array([0., 0., 0.])

    def shape_function_derivatives(self, coords=None, *args, **kwargs):
        coords = self.pointdata.x.to_numpy() if coords is None else coords
        if len(coords.shape) == 2:
            return dshp_H8_bulk(coords)
        else:
            return dshp_H8(coords)

    def volumes(self, coords=None, topo=None):
        coords = self.pointdata.x.to_numpy() if coords is None else coords
        topo = self.nodes.to_numpy() if topo is None else topo
        ecoords = cells_coords(coords, topo)
        qpos, qweight = Gauss(2, 2, 2)
        return volumes_H8(ecoords, qpos, qweight)

[((265, 296), 'numba.njit', 'njit', ([], {'nogil': '(True)', 'cache': '__cache'}), '(nogil=True, cache=__cache)\n', (269, 296), False, 'from numba import njit, prange\n'), ((412, 443), 'numba.njit', 'njit', ([], {'nogil': '(True)', 'cache': '__cache'}), '(nogil=True, cache=__cache)\n', (416, 443), False, 'from numba import njit, prange\n'), ((1546, 1592), 'numba.njit', 'njit', ([], {'nogil': '(True)', 'parallel': '(True)', 'cache': '__cache'}), '(nogil=True, parallel=True, cache=__cache)\n', (1550, 1592), False, 'from numba import njit, prange\n'), ((1842, 1873), 'numba.njit', 'njit', ([], {'nogil': '(True)', 'cache': '__cache'}), '(nogil=True, cache=__cache)\n', (1846, 1873), False, 'from numba import njit, prange\n'), ((3126, 3172), 'numba.njit', 'njit', ([], {'nogil': '(True)', 'parallel': '(True)', 'cache': '__cache'}), '(nogil=True, parallel=True, cache=__cache)\n', (3130, 3172), False, 'from numba import njit, prange\n'), ((3370, 3431), 'numba.njit', 'njit', ([], {'nogil': '(True)', 'parallel': '(True)', 'fastmath': '(True)', 'cache': '__cache'}), '(nogil=True, parallel=True, fastmath=True, cache=__cache)\n', (3374, 3431), False, 'from numba import njit, prange\n'), ((364, 418), 'numpy.array', 'np.array', (['[1, r, s, t, r * s, r * t, s * t, r * s * t]'], {}), '([1, r, s, t, r * s, r * t, s * t, r * s * t])\n', (372, 418), True, 'import numpy as np\n'), ((496, 1452), 'numpy.array', 'np.array', (['[-0.125 * r * s * t + 0.125 * r * s + 0.125 * r * t - 0.125 * r + 0.125 * s *\n    t - 0.125 * s - 0.125 * t + 0.125, 0.125 * r * s * t - 0.125 * r * s - \n    0.125 * r * t + 0.125 * r + 0.125 * s * t - 0.125 * s - 0.125 * t + \n    0.125, -0.125 * r * s * t + 0.125 * r * s - 0.125 * r * t + 0.125 * r -\n    0.125 * s * t + 0.125 * s - 0.125 * t + 0.125, 0.125 * r * s * t - \n    0.125 * r * s + 0.125 * r * t - 0.125 * r - 0.125 * s * t + 0.125 * s -\n    0.125 * t + 0.125, 0.125 * r * s * t + 0.125 * r * s - 0.125 * r * t - \n    0.125 * r - 0.125 * s * t - 0.125 * s + 0.125 * t + 0.125, -0.125 * r *\n    s * t - 0.125 * r * s + 0.125 * r * t + 0.125 * r - 0.125 * s * t - \n    0.125 * s + 0.125 * t + 0.125, 0.125 * r * s * t + 0.125 * r * s + \n    0.125 * r * t + 0.125 * r + 0.125 * s * t + 0.125 * s + 0.125 * t + \n    0.125, -0.125 * r * s * t - 0.125 * r * s - 0.125 * r * t - 0.125 * r +\n    0.125 * s * t + 0.125 * s + 0.125 * t + 0.125]'], {}), '([-0.125 * r * s * t + 0.125 * r * s + 0.125 * r * t - 0.125 * r + \n    0.125 * s * t - 0.125 * s - 0.125 * t + 0.125, 0.125 * r * s * t - \n    0.125 * r * s - 0.125 * r * t + 0.125 * r + 0.125 * s * t - 0.125 * s -\n    0.125 * t + 0.125, -0.125 * r * s * t + 0.125 * r * s - 0.125 * r * t +\n    0.125 * r - 0.125 * s * t + 0.125 * s - 0.125 * t + 0.125, 0.125 * r *\n    s * t - 0.125 * r * s + 0.125 * r * t - 0.125 * r - 0.125 * s * t + \n    0.125 * s - 0.125 * t + 0.125, 0.125 * r * s * t + 0.125 * r * s - \n    0.125 * r * t - 0.125 * r - 0.125 * s * t - 0.125 * s + 0.125 * t + \n    0.125, -0.125 * r * s * t - 0.125 * r * s + 0.125 * r * t + 0.125 * r -\n    0.125 * s * t - 0.125 * s + 0.125 * t + 0.125, 0.125 * r * s * t + \n    0.125 * r * s + 0.125 * r * t + 0.125 * r + 0.125 * s * t + 0.125 * s +\n    0.125 * t + 0.125, -0.125 * r * s * t - 0.125 * r * s - 0.125 * r * t -\n    0.125 * r + 0.125 * s * t + 0.125 * s + 0.125 * t + 0.125])\n', (504, 1452), True, 'import numpy as np\n'), ((1653, 1682), 'numpy.eye', 'np.eye', (['(3)'], {'dtype': 'pcoord.dtype'}), '(3, dtype=pcoord.dtype)\n', (1659, 1682), True, 'import numpy as np\n'), ((1718, 1755), 'numpy.zeros', 'np.zeros', (['(3, 24)'], {'dtype': 'pcoord.dtype'}), '((3, 24), dtype=pcoord.dtype)\n', (1726, 1755), True, 'import numpy as np\n'), ((1769, 1778), 'numba.prange', 'prange', (['(8)'], {}), '(8)\n', (1775, 1778), False, 'from numba import njit, prange\n'), ((1927, 3161), 'numpy.array', 'np.array', (['[[-0.125 * s * t + 0.125 * s + 0.125 * t - 0.125, -0.125 * r * t + 0.125 *\n    r + 0.125 * t - 0.125, -0.125 * r * s + 0.125 * r + 0.125 * s - 0.125],\n    [0.125 * s * t - 0.125 * s - 0.125 * t + 0.125, 0.125 * r * t - 0.125 *\n    r + 0.125 * t - 0.125, 0.125 * r * s - 0.125 * r + 0.125 * s - 0.125],\n    [-0.125 * s * t + 0.125 * s - 0.125 * t + 0.125, -0.125 * r * t + 0.125 *\n    r - 0.125 * t + 0.125, -0.125 * r * s - 0.125 * r - 0.125 * s - 0.125],\n    [0.125 * s * t - 0.125 * s + 0.125 * t - 0.125, 0.125 * r * t - 0.125 *\n    r - 0.125 * t + 0.125, 0.125 * r * s + 0.125 * r - 0.125 * s - 0.125],\n    [0.125 * s * t + 0.125 * s - 0.125 * t - 0.125, 0.125 * r * t + 0.125 *\n    r - 0.125 * t - 0.125, 0.125 * r * s - 0.125 * r - 0.125 * s + 0.125],\n    [-0.125 * s * t - 0.125 * s + 0.125 * t + 0.125, -0.125 * r * t - 0.125 *\n    r - 0.125 * t - 0.125, -0.125 * r * s + 0.125 * r - 0.125 * s + 0.125],\n    [0.125 * s * t + 0.125 * s + 0.125 * t + 0.125, 0.125 * r * t + 0.125 *\n    r + 0.125 * t + 0.125, 0.125 * r * s + 0.125 * r + 0.125 * s + 0.125],\n    [-0.125 * s * t - 0.125 * s - 0.125 * t - 0.125, -0.125 * r * t - 0.125 *\n    r + 0.125 * t + 0.125, -0.125 * r * s - 0.125 * r + 0.125 * s + 0.125]]'], {}), '([[-0.125 * s * t + 0.125 * s + 0.125 * t - 0.125, -0.125 * r * t +\n    0.125 * r + 0.125 * t - 0.125, -0.125 * r * s + 0.125 * r + 0.125 * s -\n    0.125], [0.125 * s * t - 0.125 * s - 0.125 * t + 0.125, 0.125 * r * t -\n    0.125 * r + 0.125 * t - 0.125, 0.125 * r * s - 0.125 * r + 0.125 * s - \n    0.125], [-0.125 * s * t + 0.125 * s - 0.125 * t + 0.125, -0.125 * r * t +\n    0.125 * r - 0.125 * t + 0.125, -0.125 * r * s - 0.125 * r - 0.125 * s -\n    0.125], [0.125 * s * t - 0.125 * s + 0.125 * t - 0.125, 0.125 * r * t -\n    0.125 * r - 0.125 * t + 0.125, 0.125 * r * s + 0.125 * r - 0.125 * s - \n    0.125], [0.125 * s * t + 0.125 * s - 0.125 * t - 0.125, 0.125 * r * t +\n    0.125 * r - 0.125 * t - 0.125, 0.125 * r * s - 0.125 * r - 0.125 * s + \n    0.125], [-0.125 * s * t - 0.125 * s + 0.125 * t + 0.125, -0.125 * r * t -\n    0.125 * r - 0.125 * t - 0.125, -0.125 * r * s + 0.125 * r - 0.125 * s +\n    0.125], [0.125 * s * t + 0.125 * s + 0.125 * t + 0.125, 0.125 * r * t +\n    0.125 * r + 0.125 * t + 0.125, 0.125 * r * s + 0.125 * r + 0.125 * s + \n    0.125], [-0.125 * s * t - 0.125 * s - 0.125 * t - 0.125, -0.125 * r * t -\n    0.125 * r + 0.125 * t + 0.125, -0.125 * r * s - 0.125 * r + 0.125 * s +\n    0.125]])\n', (1935, 3161), True, 'import numpy as np\n'), ((3245, 3286), 'numpy.zeros', 'np.zeros', (['(nP, 8, 3)'], {'dtype': 'pcoords.dtype'}), '((nP, 8, 3), dtype=pcoords.dtype)\n', (3253, 3286), True, 'import numpy as np\n'), ((3301, 3311), 'numba.prange', 'prange', (['nP'], {}), '(nP)\n', (3307, 3311), False, 'from numba import njit, prange\n'), ((3563, 3596), 'numpy.zeros', 'np.zeros', (['nE'], {'dtype': 'ecoords.dtype'}), '(nE, dtype=ecoords.dtype)\n', (3571, 3596), True, 'import numpy as np\n'), ((3694, 3704), 'numba.prange', 'prange', (['nE'], {}), '(nE)\n', (3700, 3704), False, 'from numba import njit, prange\n'), ((4083, 4247), 'numpy.array', 'np.array', (['[[-1.0, -1.0, -1], [1.0, -1.0, -1.0], [1.0, 1.0, -1.0], [-1.0, 1.0, -1.0],\n    [-1.0, -1.0, 1.0], [1.0, -1.0, 1.0], [1.0, 1.0, 1.0], [-1.0, 1.0, 1.0]]'], {}), '([[-1.0, -1.0, -1], [1.0, -1.0, -1.0], [1.0, 1.0, -1.0], [-1.0, 1.0,\n    -1.0], [-1.0, -1.0, 1.0], [1.0, -1.0, 1.0], [1.0, 1.0, 1.0], [-1.0, 1.0,\n    1.0]])\n', (4091, 4247), True, 'import numpy as np\n'), ((4464, 4489), 'numpy.array', 'np.array', (['[0.0, 0.0, 0.0]'], {}), '([0.0, 0.0, 0.0])\n', (4472, 4489), True, 'import numpy as np\n'), ((4963, 4989), 'dewloosh.geom.utils.cells_coords', 'cells_coords', (['coords', 'topo'], {}), '(coords, topo)\n', (4975, 4989), False, 'from dewloosh.geom.utils import cells_coords\n'), ((5014, 5028), 'dewloosh.math.numint.GaussPoints', 'Gauss', (['(2)', '(2)', '(2)'], {}), '(2, 2, 2)\n', (5019, 5028), True, 'from dewloosh.math.numint import GaussPoints as Gauss\n'), ((3763, 3781), 'numpy.linalg.det', 'np.linalg.det', (['jac'], {}), '(jac)\n', (3776, 3781), True, 'import numpy as np\n')]