# -*- coding: utf-8 -*-
"""
This file is part of FElupe.

FElupe 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.

FElupe 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 FElupe.  If not, see <http://www.gnu.org/licenses/>.
"""

import numpy as np

from ..math import cross
from ..mesh import Mesh
from ._region import Region


def boundary_cells_quad(mesh):
    "Convert the cells array of a quad mesh into a boundary cells array."

    # edges (boundary) of a quad
    i = [3, 1, 0, 2]
    j = [0, 2, 1, 3]

    cells_faces = np.dstack(
        (
            mesh.cells[:, i],
            mesh.cells[:, j],
        )
    )

    # complementary edges for the creation of "boundary" quads
    # (rotated quads with 1st edge as n-th edge of one original quad)
    t = [1, 0, 3, 2]
    k = np.array(i)[t]
    m = np.array(j)[t]

    cells = np.dstack(
        (
            cells_faces,
            mesh.cells[:, k],
            mesh.cells[:, m],
        )
    )

    return cells, cells_faces


def boundary_cells_quad8(mesh):
    "Convert the cells array of a quadratic quad mesh into a boundary cells array."

    cells_quad, cells_faces_quad = boundary_cells_quad(mesh)

    # midpoints on edges (boundary) of a quadratic quad
    m = [7, 5, 4, 6]

    cells_faces = np.dstack(
        (
            cells_faces_quad,
            mesh.cells[:, m],
        )
    )

    # complementary midpoints of edges for the creation of "boundary" quadratic quads
    # (rotated quads with 1st edge as n-th edge of one original quad)
    n = [4, 6, 5, 7]
    p = [5, 7, 6, 4]
    q = [6, 4, 7, 5]

    cells = np.dstack(
        (
            cells_quad,
            mesh.cells[:, m],
            mesh.cells[:, n],
            mesh.cells[:, p],
            mesh.cells[:, q],
        )
    )

    return cells, cells_faces


def boundary_cells_quad9(mesh):
    "Convert the cells array of a bi-quadratic quad mesh into a boundary cells array."

    cells_quad8, cells_faces_quad8 = boundary_cells_quad8(mesh)

    # midpoints on edges (boundary) of a quadratic quad
    r = [8, 8, 8, 8]

    cells = np.dstack(
        (
            cells_quad8,
            mesh.cells[:, r],
        )
    )

    return cells, cells_faces_quad8


def boundary_cells_hexahedron(mesh):
    "Convert the cells array of a hex mesh into a boundary cells array."

    # faces (boundary) of a hexahedron
    i = [0, 2, 1, 3, 0, 5]
    j = [3, 1, 0, 2, 1, 4]
    k = [7, 5, 4, 6, 2, 7]
    r = [4, 6, 5, 7, 3, 6]

    cells_faces = np.dstack(
        (
            mesh.cells[:, i],
            mesh.cells[:, j],
            mesh.cells[:, k],
            mesh.cells[:, r],
        )
    )

    # complementary faces for the creation of "boundary" hexahedrons
    # (6 rotated hexahedrons with 1st face as n-th face of
    #  one original hexahedron)
    m = [1, 3, 2, 0, 4, 1]
    n = [2, 0, 3, 1, 5, 0]
    p = [6, 4, 7, 5, 6, 3]
    q = [5, 7, 6, 4, 7, 2]

    cells = np.dstack(
        (
            cells_faces,
            mesh.cells[:, m],
            mesh.cells[:, n],
            mesh.cells[:, p],
            mesh.cells[:, q],
        )
    )

    return cells, cells_faces


def boundary_cells_hexahedron20(mesh):
    "Convert the cells array of a quadratic hex mesh into a boundary cells array."

    cells_hexahedron, cells_faces_hexahedron = boundary_cells_hexahedron(mesh)

    # midpoints on edges of faces (boundary) of a hexahedron
    i = [11, 9, 8, 10, 8, 12]
    j = [19, 17, 16, 18, 9, 15]
    k = [15, 13, 12, 14, 10, 14]
    v = [16, 18, 17, 19, 11, 13]

    cells_faces = np.dstack(
        (
            cells_faces_hexahedron,
            mesh.cells[:, i],
            mesh.cells[:, j],
            mesh.cells[:, k],
            mesh.cells[:, v],
        )
    )

    # complementary faces for the creation of "boundary" hexahedrons
    # (6 rotated hexahedrons with 1st face as n-th face of
    #  one original hexahedron)
    m = [9, 11, 10, 8, 12, 8]
    n = [18, 16, 19, 17, 13, 11]
    p = [13, 15, 14, 12, 14, 10]
    q = [17, 19, 18, 16, 15, 9]

    r = [8, 10, 9, 11, 16, 17]
    s = [10, 8, 11, 9, 17, 16]
    t = [14, 12, 15, 13, 18, 19]
    u = [12, 14, 13, 15, 19, 18]

    cells = np.dstack(
        (
            cells_hexahedron,
            mesh.cells[:, i],
            mesh.cells[:, j],
            mesh.cells[:, k],
            mesh.cells[:, v],
            mesh.cells[:, m],
            mesh.cells[:, n],
            mesh.cells[:, p],
            mesh.cells[:, q],
            mesh.cells[:, r],
            mesh.cells[:, s],
            mesh.cells[:, t],
            mesh.cells[:, u],
        )
    )

    return cells, cells_faces


def boundary_cells_hexahedron27(mesh):
    "Convert the cells array of a bi-quadratic hex mesh into a boundary cells array."

    cells_hexahedron20, cells_faces_hexahedron20 = boundary_cells_hexahedron20(mesh)

    # midpoints of faces (boundary) of a hexahedron
    i = [20, 21, 22, 23, 24, 25]

    cells_faces = np.dstack(
        (
            cells_faces_hexahedron20,
            mesh.cells[:, i],
        )
    )

    # complementary faces for the creation of "boundary" hexahedrons
    # (6 rotated hexahedrons with 1st face as n-th face of
    #  one original hexahedron)
    j = [22, 23, 21, 20, 20, 21]
    k = [23, 22, 20, 21, 21, 20]
    r = [24, 25, 24, 25, 22, 23]
    m = [25, 24, 25, 24, 23, 22]
    n = [20, 21, 22, 23, 24, 25]
    p = [21, 20, 23, 22, 25, 24]
    q = [26, 26, 26, 26, 26, 26]

    cells = np.dstack(
        (
            cells_hexahedron20,
            mesh.cells[:, j],
            mesh.cells[:, k],
            mesh.cells[:, r],
            mesh.cells[:, m],
            mesh.cells[:, n],
            mesh.cells[:, p],
            mesh.cells[:, q],
        )
    )

    return cells, cells_faces


class RegionBoundary(Region):
    r"""A numeric boundary-region as a combination of a mesh, an element and a numeric
    integration scheme (quadrature rule).

    Parameters
    ----------
    mesh : Mesh
        A mesh with points and cells.
    element : Element
        The finite element formulation to be applied on the cells.
    quadrature: Quadrature
        An element-compatible numeric integration scheme with points and weights.
    grad : bool, optional
        A flag to invoke gradient evaluation (default is True).
    only_surface: bool, optional
        A flag to use only the enclosing outline of the region (default is True).
    mask: ndarray or None, optional
        A boolean array to select a specific set of points (default is None).
    ensure_3d : bool, optional
        A flag to enforce 3d area normal vectors.

    Attributes
    ----------
    mesh : Mesh
        A mesh with points and cells.
    element : Finite element
        The finite element formulation to be applied on the cells.
    quadrature: Quadrature scheme
        An element-compatible numeric integration scheme with points and weights.
    h : ndarray
        Element shape function array ``h_aq`` of shape function ``a`` evaluated at
        quadrature point ``q``.
    dhdr : ndarray
        Partial derivative of element shape function array ``dhdr_aJq`` with shape
        function ``a`` w.r.t. natural element coordinate ``J`` evaluated at quadrature
        point ``q`` for every cell ``c`` (geometric gradient or **Jacobian**
        transformation between ``X`` and ``r``).
    dXdr : ndarray
        Geometric gradient ``dXdr_IJqc`` as partial derivative of undeformed coordinate
        ``I`` w.r.t. natural element coordinate ``J`` evaluated at quadrature point
        ``q`` for every cell ``c`` (geometric gradient or **Jacobian** transformation
        between ``X`` and ``r``).
    drdX : ndarray
        Inverse of dXdr.
    dA : ndarray
        Numeric *Differential area vectors*.
    normals : ndarray
        Area unit normal vectors.
    tangents : list of ndarray
        List of standard unit vectors.
    dV : ndarray
        Numeric *Differential volume element* as norm of *Differential area vectors*.
    dhdX : ndarray
        Partial derivative of element shape functions ``dhdX_aJqc`` of shape function
        ``a`` w.r.t. undeformed coordinate ``J`` evaluated at quadrature point ``q`` for
        every cell ``c``.

    Notes
    -----
    The gradients of the element shape functions w.r.t the undeformed coordinates are
    evaluated at all integration points of each cell in the region if the optional
    gradient argument is ``True``.

    .. math::

       \frac{\partial X_I}{\partial r_J} &= \hat{X}_{aI} \frac{
               \partial h_a}{\partial r_J
           }

       \frac{\partial h_a}{\partial X_J} &= \frac{\partial h_a}{\partial r_I}
       \frac{\partial r_I}{\partial X_J}

       dV &= \det\left(\frac{\partial X_I}{\partial r_J}\right) w

    Examples
    --------
    >>> import felupe as fem
    >>>
    >>> mesh = fem.Rectangle(n=(3, 2))
    >>> element = fem.Quad()
    >>> quadrature = fem.GaussLegendreBoundary(order=1, dim=2)
    >>>
    >>> region = fem.RegionBoundary(mesh, element, quadrature)
    >>> region
    <felupe Region object>
      Element formulation: Quad
      Quadrature rule: GaussLegendreBoundary
      Gradient evaluated: True
      Hessian evaluated: False

    The numeric differential area vectors are the products of the cofactors of the
    geometric gradient :math:`\partial X_I / \partial r_J` and the weights `w` of the
    quadrature points. The differential area vectors array is of shape
    ``(nquadraturepoints, ndim, nboundarycells)``.

    >>> region.dA
    array([[[ 0.  , -0.5 ,  0.  ,  0.5 ,  0.  ,  0.  ],
            [ 0.  , -0.5 ,  0.  ,  0.5 ,  0.  ,  0.  ]],
    <BLANKLINE>
           [[-0.25, -0.  , -0.25, -0.  ,  0.25,  0.25],
            [-0.25, -0.  , -0.25, -0.  ,  0.25,  0.25]]])

    In a boundary region, the numeric differential volumes are the magnitudes of the
    differential area vectors. For a quad mesh, the boundary cell volumes are the edge
    lengths.

    >>> region.dV.sum(axis=0)
    array([0.5, 1. , 0.5, 1. , 0.5, 0.5])

    Unit normal vectors are obtained by the ratio of the differential area vectors and
    the differential volumes.

    >>> region.dA / region.dV  ## this is equal to ``region.normals``
    array([[[ 0., -1.,  0.,  1.,  0.,  0.],
            [ 0., -1.,  0.,  1.,  0.,  0.]],
    <BLANKLINE>
           [[-1., -0., -1., -0.,  1.,  1.],
            [-1., -0., -1., -0.,  1.,  1.]]])

    The partial derivative of the first element shape function w.r.t. the undeformed
    coordinates, evaluated at the second integration point of the last element of the
    region, is obtained by:

    >>> region.dhdX[0, :, 1, -1]
    array([2.        , 0.21132487])

    The faces-cells may be used to create a mesh on the boundary.

    >>> fem.Mesh(points=mesh.points, cells=region.mesh.cells_faces, cell_type="line")
    <felupe Mesh object>
      Number of points: 6
      Number of cells:
        line: 6

    A second example shows the standard unit vectors and the area normals, located at
    the quadrature points for all edges of a quadrilateral.

    ..  pyvista-plot::
        :force_static:

        >>> import felupe as fem
        >>> import numpy as np
        >>> import pyvista as pv
        >>>
        >>> m = fem.Rectangle(n=2)
        >>> mesh = m.copy()
        >>> mesh.points[-1, 0] += 0.5
        >>> edges = fem.RegionQuadBoundary(mesh)
        >>>
        >>> start = fem.Field(edges, values=mesh.points).interpolate()
        >>> (direction,) = edges.tangents
        >>>
        >>> # 3d-data required for plotting
        >>> start = np.insert(start, len(start), 0, axis=0)
        >>> direction = np.insert(direction, len(direction), 0, axis=0)
        >>> normal = np.insert(edges.normals, len(edges.normals), 0, axis=0)
        >>>
        >>> plotter = pv.Plotter()
        ... actor = plotter.add_arrows(
        ...     start.reshape(3, -1).T,
        ...     direction.reshape(3, -1).T,
        ...     show_scalar_bar=False,
        ...     mag=1 / 7,
        ...     color="red",
        ...     label="tangents",
        ... )
        >>> actor = plotter.add_arrows(
        ...     start.reshape(3, -1).T,
        ...     normal.reshape(3, -1).T,
        ...     show_scalar_bar=False,
        ...     mag=1 / 7,
        ...     color="green",
        ...     label="normals",
        ... )
        >>> actor = plotter.add_legend()
        >>> mesh.plot(plotter=plotter, style="wireframe").show()

    A third example shows the standard unit vectors and the area normals, located at the
    quadrature points for one face of a hexahedron.

    ..  pyvista-plot::
        :force_static:

        >>> import felupe as fem
        >>> import numpy as np
        >>> import pyvista as pv
        >>>
        >>> m = fem.Cube(n=2)
        >>> mesh = m.copy()
        >>> mesh.points[-1, 0] += 0.5
        >>> faces = fem.RegionHexahedronBoundary(mesh, mask=m.x == m.x.max())
        >>>
        >>> start = fem.Field(faces, values=mesh.points).interpolate()
        >>> direction, direction_2 = faces.tangents
        >>>
        >>> plotter = pv.Plotter()
        >>> actor = plotter.add_arrows(
        ...     start.reshape(3, -1).T,
        ...     direction.reshape(3, -1).T,
        ...     show_scalar_bar=False,
        ...     mag=1 / 7,
        ...     color="red",
        ...     label="tangents (1)",
        ... )
        >>> actor = plotter.add_arrows(
        ...     start.reshape(3, -1).T,
        ...     direction_2.reshape(3, -1).T,
        ...     show_scalar_bar=False,
        ...     mag=1 / 7,
        ...     color="green",
        ...     label="tangents (2)",
        ... )
        >>> actor = plotter.add_arrows(
        ...     start.reshape(3, -1).T,
        ...     faces.normals.reshape(3, -1).T,
        ...     show_scalar_bar=False,
        ...     mag=1 / 7,
        ...     color="blue",
        ...     label="normals",
        ... )
        >>> plotter.add_legend()
        >>> mesh.plot(plotter=plotter, style="wireframe").show()

    This examples shows how to evaluate the strain on the faces of the boundary region.
    A field can't be used to evaluate the strain of a quad-cell in 3D space. However,
    it is possible to evaluate the in-plane components of the right Cauchy-Green
    deformation tensor on the boundary. With this tensor at hand, it is possible to
    evaluate any in-plane strain tensor.

    ..  pyvista-plot::
        :force_static:

        >>> import felupe as fem
        >>>
        >>> mesh = fem.Cube(n=3)
        >>> region = fem.RegionHexahedron(mesh)
        >>> field = fem.FieldContainer([fem.Field(region, dim=3)])
        >>>
        >>> field[0].values[mesh.x == 1, 0] += 0.5
        >>> field[0].values[mesh.x == 1, 2] += 0.5
        >>> field[0].values += mesh.rotate(-30, axis=1).points - mesh.points
        >>>
        >>> face = fem.RegionHexahedronBoundary(mesh, mask=mesh.x == mesh.x.max())
        >>> u = fem.FieldContainer([fem.Field(face, dim=3)])
        >>> u.link(field)
        >>>
        >>> mesh_faces = face.mesh_faces()
        >>> mesh_faces.points += field[0].values
        >>>
        >>> C = fem.math.inplane(u.evaluate.right_cauchy_green_deformation(), face.tangents)
        >>> e = fem.math.strain(u, C=C, k=0)
        >>>
        >>> view = mesh_faces.view(
        ...     cell_data={"ep": fem.math.eigvalsh(e).max(axis=0).mean(axis=-2).T}
        ... )
        >>> view.plot("ep", label="Strain (Max. Principal)", plotter=field.plot(style="wireframe")).show()
    """

    def __init__(
        self,
        mesh,
        element,
        quadrature,
        grad=True,
        only_surface=True,
        mask=None,
        ensure_3d=False,
    ):
        self.only_surface = only_surface
        self.mask = mask
        self.ensure_3d = ensure_3d

        if mesh.cell_type == "quad":
            cells, cells_faces = boundary_cells_quad(mesh)
        elif mesh.cell_type == "quad8":
            cells, cells_faces = boundary_cells_quad8(mesh)
        elif mesh.cell_type == "quad9":
            cells, cells_faces = boundary_cells_quad9(mesh)
        elif mesh.cell_type == "hexahedron":
            cells, cells_faces = boundary_cells_hexahedron(mesh)
        elif mesh.cell_type == "hexahedron20":
            cells, cells_faces = boundary_cells_hexahedron20(mesh)
        elif mesh.cell_type == "hexahedron27":
            cells, cells_faces = boundary_cells_hexahedron27(mesh)
        else:
            raise NotImplementedError("Cell type not supported.")

        cells_faces = cells_faces.reshape(-1, cells_faces.shape[-1])
        cells = cells.reshape(-1, cells.shape[-1])

        if self.only_surface:
            # sort faces, get indices of unique faces and counts
            cells_faces_sorted = np.sort(cells_faces, axis=1)
            cells_faces_unique, index, counts = np.unique(
                cells_faces_sorted, True, False, True, 0
            )

            self._index = index
            self._mask = counts == 1
        else:
            self._index = np.arange(len(cells_faces))
            self._mask = np.ones_like(self._index, dtype=bool)

        self._selection = self._index[self._mask]

        # merge with point-mask
        if mask is not None:
            point_selection = np.arange(len(mesh.points))[mask]
            self._selection = self._selection[
                np.all(
                    np.isin(cells_faces[self._selection], point_selection),
                    axis=1,
                )
            ]

        # get cell-faces and cells on boundary (unique cell-faces with one count)
        cells_on_boundary = cells[self._selection]

        # create mesh on boundary
        mesh_boundary_cells = mesh.copy()
        mesh_boundary_cells.update(cells=cells_on_boundary)
        self.mesh = mesh_boundary_cells
        self.mesh.cells_faces = cells_faces[self._selection]

        # init region and faces
        super().__init__(mesh_boundary_cells, element, quadrature, grad=grad)

        if grad:
            self.dA, self.dV, self.normals, self.tangents = self._init_faces()

    def _init_faces(self):
        "Initialize (norm of) face normals of cells."

        tangents = []

        if (
            self.mesh.cell_type == "quad"
            or self.mesh.cell_type == "quad8"
            or self.mesh.cell_type == "quad9"
        ):
            dA_1 = self.dXdr[:, 0][::-1]
            dA_1[0] = -dA_1[0]

            dX_1 = -self.dXdr[:, 0]
            dX_1[0] = -dX_1[0]

            tangents.append(dX_1 / np.linalg.norm(dX_1, axis=0))

            if self.ensure_3d:
                tangents[0] = np.insert(tangents[0], len(tangents[0]), 0, axis=0)
                other_tangent = np.zeros_like(tangents[0])
                other_tangent[2] = 1.0
                tangents.append(other_tangent)

        elif (
            self.mesh.cell_type == "hexahedron"
            or self.mesh.cell_type == "hexahedron20"
            or self.mesh.cell_type == "hexahedron27"
        ):
            dA_1 = cross(self.dXdr[:, 0], self.dXdr[:, 1])

            tangents.append(-self.dXdr[:, 0] / np.linalg.norm(self.dXdr[:, 0], axis=0))
            tangents.append(self.dXdr[:, 1] / np.linalg.norm(self.dXdr[:, 1], axis=0))

        dA = -dA_1 * self.quadrature.weights.reshape(-1, 1)

        # norm and unit normal vector
        dV = np.linalg.norm(dA, axis=0)
        normals = dA / dV

        if self.ensure_3d:
            if dA.shape[0] == 2:
                dA = np.pad(dA, ((0, 1), (0, 0), (0, 0)))
                normals = np.pad(normals, ((0, 1), (0, 0), (0, 0)))

        return dA, dV, normals, tangents

    def mesh_faces(self):
        "Return a Mesh with face-cells on the selected boundary."

        face_type = {
            "quad": "line",
            "hexahedron": "quad",
            "quad8": "line3",
            "quad9": "line3",
        }[self.mesh.cell_type]
        return Mesh(self.mesh.points, self.mesh.cells_faces, face_type)
