# Copyright (C) 2024 Garth N. Wells and Paul T. Kühner
#
# This file is part of DOLFINx (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later
"""Finite elements."""

from functools import singledispatch

import numpy as np
import numpy.typing as npt

import basix
import basix.ufl
from dolfinx import cpp as _cpp


class CoordinateElement:
    """Coordinate element describing the geometry map for mesh cells."""

    _cpp_object: _cpp.fem.CoordinateElement_float32 | _cpp.fem.CoordinateElement_float64

    def __init__(
        self, cmap: _cpp.fem.CoordinateElement_float32 | _cpp.fem.CoordinateElement_float64
    ):
        """Create a coordinate map element.

        Note:
            This initialiser is for internal use and not normally called
            in user code. Use :func:`coordinate_element` to create a
            CoordinateElement.

        Args:
            cmap: A C++ CoordinateElement.
        """
        assert isinstance(
            cmap, _cpp.fem.CoordinateElement_float32 | _cpp.fem.CoordinateElement_float64
        )
        self._cpp_object = cmap

    @property
    def dtype(self) -> np.dtype:
        """Scalar type for the coordinate element."""
        return np.dtype(self._cpp_object.dtype)

    @property
    def dim(self) -> int:
        """Dimension of the coordinate element space.

        This is number of basis functions that span the coordinate
        space, e.g., for a linear triangle cell the dimension will be 3.
        """
        return self._cpp_object.dim

    def hash(self) -> int:
        """Hash identifier of the coordinate element."""
        return self._cpp_object.hash()

    def create_dof_layout(self) -> _cpp.fem.ElementDofLayout:
        """Compute and return the dof layout"""
        return self._cpp_object.create_dof_layout()

    def push_forward(
        self,
        X: npt.NDArray[np.float32] | npt.NDArray[np.float64],
        cell_geometry: npt.NDArray[np.float32] | npt.NDArray[np.float64],
    ) -> npt.NDArray[np.float32] | npt.NDArray[np.float64]:
        """Push points on the reference cell forward to the physical cell.

        Args:
            X: Coordinates of points on the reference cell,
                ``shape=(num_points, topological_dimension)``.
            cell_geometry: Coordinate 'degrees-of-freedom' (nodes) of
                the cell, ``shape=(num_geometry_basis_functions,
                geometrical_dimension)``. Can be created by accessing
                ``geometry.x[geometry.dofmap.cell_dofs(i)]``,

        Returns:
            Physical coordinates of the points reference points ``X``.
        """
        return self._cpp_object.push_forward(X, cell_geometry)

    def pull_back(
        self,
        x: npt.NDArray[np.float32] | npt.NDArray[np.float64],
        cell_geometry: npt.NDArray[np.float32] | npt.NDArray[np.float64],
    ) -> npt.NDArray[np.float32] | npt.NDArray[np.float64]:
        """Pull points on the physical cell back to the reference cell.

        For non-affine cells, the pull-back is a nonlinear operation.

        Args:
            x: Physical coordinates to pull back to the reference cells,
                ``shape=(num_points, geometrical_dimension)``.
            cell_geometry: Physical coordinates describing the cell,
                shape ``(num_of_geometry_basis_functions,
                geometrical_dimension)``. They can be created by accessing
                ``geometry.x[geometry.dofmap.cell_dofs(i)]``,

        Returns:
            Reference coordinates of the physical points ``x``.
        """
        return self._cpp_object.pull_back(x, cell_geometry)

    @property
    def variant(self) -> int:
        """Lagrange variant of the coordinate element.

        Note:
            The return type is an integer. A Basix enum can be
            created using ``basix.LagrangeVariant(value)``.
        """
        return self._cpp_object.variant

    @property
    def degree(self) -> int:
        """Polynomial degree of the coordinate element."""
        return self._cpp_object.degree


@singledispatch
def coordinate_element(
    celltype: _cpp.mesh.CellType,
    degree: int,
    variant=int(basix.LagrangeVariant.unset),
    dtype: npt.DTypeLike = np.float64,
):
    """Create a Lagrange CoordinateElement from element metadata.

    Coordinate elements are typically used to create meshes.

    Args:
        celltype: Cell shape
        degree: Polynomial degree of the coordinate element map.
        variant: Basix Lagrange variant (affects node placement).
        dtype: Scalar type for the coordinate element.

    Returns:
        A coordinate element.
    """
    if np.issubdtype(dtype, np.float32):
        return CoordinateElement(_cpp.fem.CoordinateElement_float32(celltype, degree, variant))
    elif np.issubdtype(dtype, np.float64):
        return CoordinateElement(_cpp.fem.CoordinateElement_float64(celltype, degree, variant))
    else:
        raise RuntimeError("Unsupported dtype.")


@coordinate_element.register(basix.finite_element.FiniteElement)
def _(e: basix.finite_element.FiniteElement):
    """Create a Lagrange CoordinateElement from a Basix finite element.

    Coordinate elements are typically used when creating meshes.

    Args:
        e: Basix finite element.

    Returns:
        A coordinate element.
    """
    try:
        return CoordinateElement(_cpp.fem.CoordinateElement_float32(e._e))
    except TypeError:
        return CoordinateElement(_cpp.fem.CoordinateElement_float64(e._e))


class FiniteElement:
    _cpp_object: _cpp.fem.FiniteElement_float32 | _cpp.fem.FiniteElement_float64

    def __init__(
        self,
        cpp_object: _cpp.fem.FiniteElement_float32 | _cpp.fem.FiniteElement_float64,
    ):
        """Creates a Python wrapper for the exported finite element class.

        Note:
            Do not use this constructor directly. Instead use
            :func:`finiteelement`.

        Args:
            The underlying cpp instance that this object will wrap.
        """
        self._cpp_object = cpp_object

    def __eq__(self, other):
        return self._cpp_object == other._cpp_object

    @property
    def dtype(self) -> np.dtype:
        """Geometry type of the Mesh that the FunctionSpace is defined
        on."""
        return self._cpp_object.dtype

    @property
    def basix_element(self) -> basix.finite_element.FiniteElement:
        """Return underlying Basix C++ element (if it exists).

        Raises:
            Runtime error if Basix element does not exist.
        """
        return self._cpp_object.basix_element

    @property
    def num_sub_elements(self) -> int:
        """Number of sub elements (for a mixed or blocked element)."""
        return self._cpp_object.num_sub_elements

    @property
    def value_shape(self) -> npt.NDArray[np.integer]:
        """Value shape of the finite element field.

        The value shape describes the shape of the finite element field,
        e.g. ``{}`` for a scalar, ``{2}`` for a vector in 2D, ``{3, 3}``
        for a rank-2 tensor in 3D, etc.
        """
        return self._cpp_object.value_shape

    @property
    def interpolation_points(self) -> npt.NDArray[np.floating]:
        """Points on the reference cell at which an expression needs to be
        evaluated in order to interpolate the expression in the finite
        element space.

        Interpolation point coordinates on the reference cell, returning
        the coordinates data (row-major) storage with shape
        ``(num_points, tdim)``.

        Note:
            For Lagrange elements the points will just be the nodal
            positions. For other elements the points will typically be the
            quadrature points used to evaluate moment degrees of freedom.
        """
        return self._cpp_object.interpolation_points()

    @property
    def interpolation_ident(self) -> bool:
        """Check if interpolation into the finite element space is an
        identity operation given the evaluation on an expression at
        specific points, i.e. the degree-of-freedom are equal to point
        evaluations. The function will return `true` for Lagrange
        elements."""
        return self._cpp_object.interpolation_ident

    @property
    def space_dimension(self) -> int:
        """Dimension of the finite element function space (the number of
        degrees-of-freedom for the element).

        For 'blocked' elements, this function returns the dimension of the
        full element rather than the dimension of the base element.
        """
        return self._cpp_object.space_dimension

    @property
    def needs_dof_transformations(self) -> bool:
        """Check if DOF transformations are needed for this element.

        DOF transformations will be needed for elements which might not be
        continuous when two neighbouring cells disagree on the orientation
        of a shared sub-entity, and when this cannot be corrected for by
        permuting the DOF numbering in the dofmap.

        For example, Raviart-Thomas elements will need DOF transformations,
        as the neighbouring cells may disagree on the orientation of a
        basis function, and this orientation cannot be corrected for by
        permuting the DOF numbers on each cell.
        """
        return self._cpp_object.needs_dof_transformations

    @property
    def signature(self) -> str:
        """String identifying the finite element."""
        return self._cpp_object.signature

    def T_apply(
        self, x: npt.NDArray[np.floating], cell_permutations: npt.NDArray[np.uint32], dim: int
    ) -> None:
        """Transform basis functions from the reference element ordering
        and orientation to the globally consistent physical element
        ordering and orientation.

        Args:
            x: Data to transform (in place). The shape is
                ``(num_cells, n, dim)``, where ``n`` is the number degrees-
                of-freedom and the data is flattened (row-major).
            cell_permutations: Permutation data for the cell.
            dim: Number of columns in ``data``.

        Note:
            Exposed for testing. Function is not vectorised across multiple
            cells. Please see `basix.numba_helpers` for performant
            versions.
        """
        self._cpp_object.T_apply(x, cell_permutations, dim)

    def Tt_apply(
        self, x: npt.NDArray[np.floating], cell_permutations: npt.NDArray[np.uint32], dim: int
    ) -> None:
        """Apply the transpose of the operator applied by T_apply().

        Args:
            x: Data to transform (in place). The shape is
                ``(num_cells, n, dim)``, where ``n`` is the number degrees-
                of-freedom and the data is flattened (row-major).
            cell_permutations: Permutation data for the cells
            dim: Number of columns in ``data``.
        """
        self._cpp_object.Tt_apply(x, cell_permutations, dim)

    def Tt_inv_apply(
        self, x: npt.NDArray[np.floating], cell_permutations: npt.NDArray[np.uint32], dim: int
    ) -> None:
        """Apply the inverse transpose of the operator applied by
        T_apply().

        Args:
            x: Data to transform (in place). The shape is
                ``(num_cells, n, dim)``, where ``n`` is the number degrees-
                of-freedom and the data is flattened (row-major).
            cell_permutations: Permutation data for the cells
            dim: Number of columns in ``data``.
        """
        self._cpp_object.Tt_inv_apply(x, cell_permutations, dim)


def finiteelement(
    cell_type: _cpp.mesh.CellType,
    ufl_e: basix.ufl._ElementBase,
    FiniteElement_dtype: np.dtype,
) -> FiniteElement:
    """Create a DOLFINx element from a basix.ufl element.

    Args:
        cell_type: Element cell type, see ``mesh.CellType``
        ufl_e: UFL element, holding quadrature rule and other properties of
            the selected element.
        FiniteElement_dtype: Geometry type of the element.
    """
    if np.issubdtype(FiniteElement_dtype, np.float32):
        CppElement = _cpp.fem.FiniteElement_float32
    elif np.issubdtype(FiniteElement_dtype, np.float64):
        CppElement = _cpp.fem.FiniteElement_float64
    else:
        raise ValueError(f"Unsupported dtype: {FiniteElement_dtype}")

    if ufl_e.is_mixed:
        elements = [
            finiteelement(cell_type, e, FiniteElement_dtype)._cpp_object  # type: ignore
            for e in ufl_e.sub_elements
        ]
        return FiniteElement(CppElement(elements))
    elif ufl_e.is_quadrature:
        return FiniteElement(
            CppElement(
                cell_type,
                ufl_e.custom_quadrature()[0],
                ufl_e.reference_value_shape,
                ufl_e.is_symmetric,
            )
        )
    else:
        basix_e = ufl_e.basix_element._e
        value_shape = ufl_e.reference_value_shape if ufl_e.block_size > 1 else None
        return FiniteElement(CppElement(basix_e, value_shape, ufl_e.is_symmetric))