# -*- 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)

[(11, 'numba.njit', 'njit', False, 'from numba import njit, prange\n'), (17, 'numba.njit', 'njit', False, 'from numba import njit, prange\n'), (39, 'numba.njit', 'njit', False, 'from numba import njit, prange\n'), (49, 'numba.njit', 'njit', False, 'from numba import njit, prange\n'), (80, 'numba.njit', 'njit', False, 'from numba import njit, prange\n'), (89, 'numba.njit', 'njit', False, 'from numba import njit, prange\n'), (14, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (20, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (41, 'numpy.eye', 'np.eye', True, 'import numpy as np\n'), (43, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (44, 'numba.prange', 'prange', False, 'from numba import njit, prange\n'), (52, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (83, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (84, 'numba.prange', 'prange', False, 'from numba import njit, prange\n'), (93, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (97, 'numba.prange', 'prange', False, 'from numba import njit, prange\n'), (122, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (133, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (145, 'dewloosh.geom.utils.cells_coords', 'cells_coords', False, 'from dewloosh.geom.utils import cells_coords\n'), (146, 'dewloosh.math.numint.GaussPoints', 'Gauss', True, 'from dewloosh.math.numint import GaussPoints as Gauss\n'), (99, 'numpy.linalg.det', 'np.linalg.det', True, 'import numpy as np\n')]