# Copyright (C) 2009-2026 Chris N. Richardson, Garth N. Wells,
# Michal Habera and Jørgen S. Dokken
#
# This file is part of DOLFINx (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later
"""Finite element function spaces and functions."""

from __future__ import annotations

import typing
from collections.abc import Callable, Sequence
from functools import cached_property, singledispatch
from typing import Generic

import numpy as np
import numpy.typing as npt

import basix
import ufl
from dolfinx import cpp as _cpp
from dolfinx import default_scalar_type, jit, la
from dolfinx.fem.dofmap import DofMap
from dolfinx.fem.element import FiniteElement, finiteelement
from dolfinx.geometry import PointOwnershipData
from dolfinx.typing import Real, Scalar

if typing.TYPE_CHECKING:
    from mpi4py import MPI as _MPI

    from dolfinx.mesh import EntityMap as _EntityMap
    from dolfinx.mesh import Mesh


class Constant(ufl.Constant, Generic[Scalar]):
    """A constant with respect to a domain."""

    _cpp_object: (
        _cpp.fem.Constant_complex64
        | _cpp.fem.Constant_complex128
        | _cpp.fem.Constant_float32
        | _cpp.fem.Constant_float64
    )

    def __init__(
        self,
        domain,
        c: float | np.floating | complex | np.complexfloating | Sequence | np.ndarray,
    ):
        """Create a Constant.

        Args:
            domain: DOLFINx or UFL mesh
            c: Value of the constant.
        """
        c = np.asarray(c)
        super().__init__(domain, c.shape)
        try:
            if np.issubdtype(c.dtype, np.complex64):
                self._cpp_object = _cpp.fem.Constant_complex64(c)
            elif np.issubdtype(c.dtype, np.complex128):
                self._cpp_object = _cpp.fem.Constant_complex128(c)
            elif np.issubdtype(c.dtype, np.float32):
                self._cpp_object = _cpp.fem.Constant_float32(c)
            elif np.issubdtype(c.dtype, np.float64):
                self._cpp_object = _cpp.fem.Constant_float64(c)
            else:
                raise RuntimeError("Unsupported dtype")
        except AttributeError:
            raise AttributeError("Constant value must have a dtype attribute.")

    @property
    def value(self):
        """The value of the constant."""
        return self._cpp_object.value

    @value.setter
    def value(self, v: npt.NDArray[Scalar]) -> None:
        np.copyto(self._cpp_object.value, np.asarray(v))

    @property
    def dtype(self) -> npt.DTypeLike:
        """Value dtype of the constant."""
        return np.dtype(self._cpp_object.dtype)

    def __float__(self) -> float:
        """Real representation of the constant."""
        if self.ufl_shape or self.ufl_free_indices:
            raise TypeError("Cannot evaluate a nonscalar expression to a scalar value.")

        return float(self.value)

    def __complex__(self) -> complex:
        """Complex representation of the constant."""
        if self.ufl_shape or self.ufl_free_indices:
            raise TypeError("Cannot evaluate a nonscalar expression to a scalar value.")

        return complex(self.value)


class Expression(Generic[Scalar]):
    """An object for evaluating functions of finite element functions.

    Represents a mathematical expression evaluated at a pre-defined set
    of points on the reference cell. This class closely follows the
    concept of a UFC Expression.

    This functionality can be used to evaluate a gradient of a Function
    at the quadrature points in all cells. This evaluated gradient can
    then be used as input to a non-FEniCS function that calculates a
    material constitutive model.

    """

    _ufl_expression: ufl.core.expr.Expr
    _argument_space: FunctionSpace | None
    _cpp_object: (
        _cpp.fem.Expression_complex64
        | _cpp.fem.Expression_complex128
        | _cpp.fem.Expression_float32
        | _cpp.fem.Expression_float64
    )
    _code: str

    def __init__(
        self,
        e: ufl.core.expr.Expr,
        X: np.ndarray,
        comm: _MPI.Comm | None = None,
        form_compiler_options: dict | None = None,
        jit_options: dict | None = None,
        dtype: npt.DTypeLike | None = None,
        entity_maps: list[_EntityMap] | None = None,
    ):
        """Create an Expression.

        Args:
            e: UFL expression.
            X: Array of points of shape ``(num_points, tdim)`` on the
                reference element.
            comm: Communicator that the Expression is defined on.
            form_compiler_options: Options used in FFCx compilation of
                this Expression. Run ``ffcx --help`` in the commandline
                to see all available options.
            jit_options: Options controlling JIT compilation of C code.
            dtype: Type of the Expression values. If ``None``,
                the dtype is deduced from the UFL expression.
            entity_maps: Maps between different meshes.

        Note:
            This wrapper is responsible for the FFCx compilation of the
            UFL Expr and attaching the correct data to the underlying
            C++ Expression.
        """
        assert X.ndim < 3
        num_points = X.shape[0] if X.ndim == 2 else 1
        _X = np.reshape(X, (num_points, -1))

        # Get MPI communicator
        if comm is None:
            try:
                domains = ufl.domain.extract_domains(e)
                if len(domains) == 0:
                    raise RuntimeError("Could not extract MPI communicator for Expression.")
                domain = domains[0]
                assert isinstance(domain, ufl.Mesh)
                mesh = domain.ufl_cargo()
                comm = mesh.comm
            except AttributeError:
                print(
                    "Could not extract MPI communicator for Expression. "
                    + "Maybe you need to pass a communicator?"
                )
                raise

        # Attempt to deduce dtype
        if dtype is None:
            dtype = getattr(e, "dtype", default_scalar_type)

        # Compile UFL expression with JIT
        if form_compiler_options is None:
            form_compiler_options = dict()
        form_compiler_options["scalar_type"] = dtype
        self._ufcx_expression, module, self._code = jit.ffcx_jit(
            comm, (e, _X), form_compiler_options=form_compiler_options, jit_options=jit_options
        )
        self._ufl_expression = e

        # Prepare coefficients data. For every coefficient in expression
        # take its C++ object.
        original_coefficients = ufl.algorithms.extract_coefficients(e)
        coeffs = [
            original_coefficients[
                self._ufcx_expression.original_coefficient_positions[i]
            ]._cpp_object
            for i in range(self._ufcx_expression.num_coefficients)
        ]
        ufl_constants = ufl.algorithms.analysis.extract_constants(e)
        constants = [constant._cpp_object for constant in ufl_constants]
        arguments = ufl.algorithms.extract_arguments(e)
        if len(arguments) == 0:
            self._argument_space = None
        elif len(arguments) == 1:
            self._argument_space = arguments[0].ufl_function_space()._cpp_object
        else:
            raise RuntimeError("Expressions with more that one Argument not allowed.")

        def _create_expression(dtype):
            if np.issubdtype(dtype, np.float32):
                return _cpp.fem.create_expression_float32
            elif np.issubdtype(dtype, np.float64):
                return _cpp.fem.create_expression_float64
            elif np.issubdtype(dtype, np.complex64):
                return _cpp.fem.create_expression_complex64
            elif np.issubdtype(dtype, np.complex128):
                return _cpp.fem.create_expression_complex128
            else:
                raise NotImplementedError(f"Type {dtype} not supported.")

        _entity_maps = (
            [entity_map._cpp_object for entity_map in entity_maps]
            if entity_maps is not None
            else []
        )
        ffi = module.ffi
        self._cpp_object = _create_expression(dtype)(
            ffi.cast("uintptr_t", ffi.addressof(self._ufcx_expression)),
            coeffs,
            constants,
            _entity_maps,
            self.argument_space,
        )

    def eval(
        self,
        mesh: Mesh,
        entities: npt.NDArray[np.int32],
        values: npt.NDArray[Scalar] | None = None,
    ) -> npt.NDArray[Scalar]:
        """Evaluate Expression on entities.

        Args:
            mesh: Mesh to evaluate Expression on.
            entities: Entities to evaluate the Expression over. For
                cells, it is a list of cell indices. For facets, it is a
                2D array of (cell index, local facet index).
            values: Array to fill with evaluated values. If ``None``,
                storage will be allocated. Otherwise it must have shape
                ``(entities.shape[0], num_points, *value_shape)`` if
                there is not argument function and shape
                ``(entities.shape[0], num_points, *value_shape,
                argument_space_dim)`` if the Expression does have an
                argument function.

        Returns:
            Expression evaluated at points for ``entities``. Shape is
            ``(entities.shape[0], num_points, *value_shape)`` if there
            is no argument function, or shape is ``(entities.shape[0],
            num_points, *value_shape, argument_space_dim)`` if the
            Expression does have an argument function.
        """
        _entities = np.asarray(entities, dtype=np.int32)
        if (tdim := mesh.topology.dim) != (expr_dim := self._cpp_object.X().shape[1]):
            assert expr_dim == tdim - 1
            assert entities.ndim == 2, (
                "entities list should have two dimensions for expression evaluation on facets."
            )

        if self.argument_space is None:
            values_shape = (_entities.shape[0], self.X().shape[0], *self.value_shape)
        else:
            values_shape = (
                _entities.shape[0],
                self.X().shape[0],
                *self.value_shape,
                self.argument_space.element.space_dimension,
            )

        # Allocate memory for result if u was not provided
        if values is None:
            values = np.zeros(values_shape, dtype=self.dtype)
        else:
            if values.shape != values_shape:
                raise TypeError("Passed values array does not have correct shape.")
            if values.dtype != self.dtype:
                raise TypeError("Passed values array does not have correct dtype.")

        constants = _cpp.fem.pack_constants(self._cpp_object)
        coeffs = _cpp.fem.pack_coefficients(self._cpp_object, mesh._cpp_object, _entities)
        _cpp.fem.tabulate_expression(
            values, self._cpp_object, constants, coeffs, mesh._cpp_object, _entities
        )
        return values

    def X(self) -> npt.NDArray:
        """Evaluation points on the reference cell."""
        return self._cpp_object.X()

    @property
    def ufl_expression(self) -> ufl.core.expr.Expr:
        """Original UFL Expression."""
        return self._ufl_expression

    @property
    def value_shape(self) -> np.ndarray:
        """Value shape of the Expression."""
        return self._cpp_object.value_shape

    @property
    def value_size(self) -> int:
        """Value size of the expression."""
        return self._cpp_object.value_size

    @property
    def argument_space(self) -> FunctionSpace | None:
        """Argument function space if Expression has argument."""
        return self._argument_space

    @property
    def ufcx_expression(self):
        """The compiled ufcx_expression object."""
        return self._ufcx_expression

    @property
    def code(self) -> str:
        """C code strings."""
        return self._code

    @property
    def dtype(self) -> npt.DTypeLike:
        """Expression value dtype."""
        return np.dtype(self._cpp_object.dtype)


class Function(ufl.Coefficient, Generic[Scalar]):
    """A finite element function.

    A finite element function is represented by a function space
    (domain, element and dofmap) and a vector holding the
    degrees-of-freedom.

    """

    _cpp_object: (
        _cpp.fem.Function_complex64
        | _cpp.fem.Function_complex128
        | _cpp.fem.Function_float32
        | _cpp.fem.Function_float64
    )

    _x: la.Vector[Scalar]

    def __init__(
        self,
        V: FunctionSpace,
        x: la.Vector[Scalar] | None = None,
        name: str | None = None,
        dtype: npt.DTypeLike | None = None,
    ):
        """Initialize a finite element Function.

        Args:
            V: The function space that the Function is defined on.
            x: Function degree-of-freedom vector. Typically required
                only when reading a saved Function from file.
            name: Function name.
            dtype: Scalar type. Is not set, the DOLFINx default scalar
                type is used.
        """
        if x is not None:
            if dtype is None:
                dtype = x.array.dtype
            else:
                assert x.array.dtype == dtype, "Incompatible Vector and dtype."
        else:
            if dtype is None:
                dtype = default_scalar_type

        assert np.issubdtype(V.element.dtype, np.dtype(dtype).type(0).real.dtype), (
            "Incompatible FunctionSpace dtype and requested dtype."
        )

        # Create cpp Function
        def functiontype(dtype):
            if np.issubdtype(dtype, np.float32):
                return _cpp.fem.Function_float32
            elif np.issubdtype(dtype, np.float64):
                return _cpp.fem.Function_float64
            elif np.issubdtype(dtype, np.complex64):
                return _cpp.fem.Function_complex64
            elif np.issubdtype(dtype, np.complex128):
                return _cpp.fem.Function_complex128
            else:
                raise NotImplementedError(f"Type {dtype} not supported.")

        if x is not None:
            self._cpp_object = functiontype(dtype)(V._cpp_object, x._cpp_object)  # type: ignore
        else:
            self._cpp_object = functiontype(dtype)(V._cpp_object)  # type: ignore

        # Initialize the ufl.FunctionSpace
        super().__init__(V.ufl_function_space())

        # Set name
        if name is None:
            self.name = "f"
        else:
            self.name = name

        # Store DOLFINx FunctionSpace object
        self._V = V

        # Store Python wrapper around the underlying Vector
        self._x = la.Vector(self._cpp_object.x)

    @property
    def function_space(self) -> FunctionSpace:
        """FunctionSpace that the Function is defined on."""
        return self._V

    def eval(
        self,
        x: npt.ArrayLike,
        cells: npt.NDArray[np.int32],
        u: None | npt.NDArray[Scalar] = None,
        tol: float = 1.0e-6,
        maxit: int = 15,
    ) -> npt.NDArray[Scalar]:
        """Evaluate Function at points x.

        Args:
            x: Points with shape (num_points, 3)
            cells: Array with cell indices, with shape (num_points,), where
              cell[i] is the index of the cell containing point x[i].
              If the cell index is negative the point is ignored.
            u: Array to put evaluated data in.
            tol: Tolerance for convergence in Newton method for
                nonaffine pullbacks.
            maxit: Maximum number of Newton iterations for
                nonaffine pullbacks.
        """
        # Make sure input coordinates are a NumPy array
        _x = np.asarray(x, dtype=self._V.mesh.geometry.x.dtype)
        assert _x.ndim < 3
        if len(_x) == 0:
            _x = np.zeros((0, 3), dtype=self._V.mesh.geometry.x.dtype)
        else:
            shape0 = _x.shape[0] if _x.ndim == 2 else 1
            _x = np.reshape(_x, (shape0, -1))
        num_points = _x.shape[0]
        if _x.shape[1] != 3:
            raise ValueError("Coordinate(s) for Function evaluation must have length 3.")

        # Make sure cells are a NumPy array
        _cells = np.asarray(cells, dtype=np.int32)
        assert _cells.ndim < 2
        num_points_c = _cells.shape[0] if _cells.ndim == 1 else 1
        _cells = np.reshape(_cells, num_points_c)

        # Allocate memory for return value if not provided
        if u is None:
            value_size = self._V.value_size
            u = np.empty((num_points, value_size), self.dtype)

        self._cpp_object.eval(_x, _cells, u, tol, maxit)  # type: ignore
        if num_points == 1:
            u = np.reshape(u, (-1,))
        return u

    def interpolate_nonmatching(
        self,
        u0: Function[Scalar],
        cells: npt.NDArray[np.int32],
        interpolation_data: PointOwnershipData,
        tol: float = 1e-6,
        maxit: int = 15,
    ) -> None:
        """Interpolate a Function on a non-matching mesh.

        Args:
            u0: Function to interpolate.
            cells: The cells to interpolate over. If ``None`` then all
                cells are interpolated over.
            interpolation_data: Data needed to interpolate functions
                defined on other meshes. Created by
                :func:`dolfinx.fem.create_interpolation_data`.
            tol: Tolerance for convergence in Newton method for nonaffine
                pullbacks. Ignored if mesh geometry is affine.
            maxit: Maximum number of iterations for nonaffine pullback.
                Ignored if mesh geometry is affine.
        """
        self._cpp_object.interpolate(
            u0._cpp_object, cells, tol, maxit, interpolation_data._cpp_object
        )  # type: ignore

    def interpolate(
        self,
        u0: Callable | Expression[Scalar] | Function[Scalar],
        cells0: npt.NDArray[np.int32] | None = None,
        cells1: npt.NDArray[np.int32] | None = None,
    ) -> None:
        """Interpolate an expression.

        Args:
            u0: Callable function, Expression or Function to
               interpolate.
            cells0: Cells in mesh associated with ``u0`` to interpolate
                over. If ``None`` then all cells are interpolated over.
            cells1: Cells in the mesh associated with ``self`` to
                interpolate over. If ``None``, then taken to be the same
                cells as ``cells0``. If ``cells1`` is not ``None`` it
                must have the same length as ``cells0``.
        """

        @singledispatch
        def _interpolate(u0):
            """Interpolate a cpp.fem.Function."""
            self._cpp_object.interpolate(u0, cells0, cells1)

        @_interpolate.register(Function)
        def _(u0: Function):
            """Interpolate a fem.Function."""
            self._cpp_object.interpolate(u0._cpp_object, cells0, cells1)

        @_interpolate.register(int)
        def _(u0_ptr: int):
            """Interpolate using a pointer to a function f(x)."""
            self._cpp_object.interpolate_ptr(u0_ptr, cells0)

        @_interpolate.register(Expression)
        def _(e0: Expression):
            """Interpolate a fem.Expression."""
            self._cpp_object.interpolate_expr(e0._cpp_object, cells0, cells1)

        try:
            # u is a Function or Expression (or pointer to one)
            _interpolate(u0)
        except TypeError:
            # u0 is callable
            assert callable(u0)
            x = _cpp.fem.interpolation_coords(
                self._V.element._cpp_object, self._V.mesh.geometry._cpp_object, cells0
            )
            self._cpp_object.interpolate_f(np.asarray(u0(x), dtype=self.dtype), cells0)

    def copy(self) -> Function[Scalar]:
        """Create a copy of the Function.

        The function space is shared and the degree-of-freedom vector is
        copied.

        Returns:
            A new Function with a copy of the degree-of-freedom vector.
        """
        return Function(
            self.function_space,
            la.Vector(type(self.x._cpp_object)(self.x._cpp_object)),
            name=self.name,
        )

    @property
    def x(self) -> la.Vector[Scalar]:
        """Vector holding the degrees-of-freedom."""
        return self._x

    @property
    def dtype(self) -> npt.DTypeLike:
        """Function value dtype."""
        return np.dtype(self._cpp_object.x.array.dtype)

    @property
    def name(self) -> str:
        """Name of the Function."""
        return self._cpp_object.name  # type: ignore

    @name.setter
    def name(self, name):
        self._cpp_object.name = name

    def __str__(self) -> str:
        """Pretty print representation."""
        return self.name

    def sub(self, i: int) -> Function[Scalar]:
        """Return a sub-function (a view into the ``Function``).

        Sub-functions are indexed ``i = 0, ..., N-1``, where ``N`` is
        the number of sub-spaces.

        Args:
            i: Index of the sub-function to extract.

        Returns:
            A view into the parent ``Function``.

        Note:
            If the sub-Function is re-used, for performance reasons the
            returned ``Function`` should be stored by the caller to
            avoid repeated re-computation of the subspace.
        """
        return Function(self._V.sub(i), self.x, name=f"{self!s}_{i}")

    def split(self) -> tuple[Function[Scalar], ...]:
        """Extract (any) sub-functions.

        A sub-function can be extracted from a discrete function that is
        in a mixed, vector, or tensor FunctionSpace. The sub-function
        resides in the subspace of the mixed space.

        Returns:
            First level of subspaces of the function space.
        """
        num_sub_spaces = self.function_space.num_sub_spaces
        if num_sub_spaces == 1:
            raise RuntimeError("No subfunctions to extract")
        return tuple(self.sub(i) for i in range(num_sub_spaces))

    def collapse(self) -> Function[Scalar]:
        """Create a collapsed version of this Function."""
        u_collapsed = self._cpp_object.collapse()  # type: ignore
        V_collapsed = FunctionSpace(
            self.function_space._mesh,
            self.ufl_element(),  # type: ignore
            u_collapsed.function_space,
        )
        return Function(V_collapsed, la.Vector(u_collapsed.x))


class ElementMetaData(typing.NamedTuple):
    """Data for representing a finite element.

    :param family: Element type.
    :param degree: Polynomial degree of the element.
    :param shape: Shape for vector/tensor valued elements that are
        constructed from blocked scalar elements (e.g., Lagrange).
    :param symmetry: Symmetry option for blocked tensor elements.
    """

    family: str
    degree: int
    shape: tuple[int, ...] | None = None
    symmetry: bool | None = None


def functionspace(
    mesh: Mesh,
    element: (
        ufl.finiteelement.AbstractFiniteElement
        | ElementMetaData
        | tuple[str, int]
        | tuple[str, int, tuple]
        | tuple[str, int, tuple, bool]
    ),
) -> FunctionSpace:
    """Create a finite element function space.

    Args:
        mesh: Mesh that space is defined on.
        element: Finite element description.

    Returns:
        A function space.
    """
    # Create UFL element
    dtype = mesh.geometry.x.dtype
    try:
        e = ElementMetaData(*element)  # type: ignore
        ufl_e = basix.ufl.element(
            e.family,
            mesh.basix_cell(),  # type: ignore
            e.degree,
            shape=e.shape,
            symmetry=e.symmetry,
            dtype=dtype,
        )
    except TypeError:
        ufl_e = element  # type: ignore

    # Check that element and mesh cell types match
    if ((domain := mesh.ufl_domain()) is None) or ufl_e.cell != domain.ufl_cell():
        raise ValueError("Non-matching UFL cell and mesh cell shapes.")
    # Create DOLFINx objects
    element = finiteelement(mesh.topology.cell_type, ufl_e, dtype)  # type: ignore

    if ufl_e.is_real:
        cpp_dofmap = _cpp.fem.build_real_element_dofmap(
            mesh.topology._cpp_object,
            element.basix_element.entity_dofs,  # type: ignore
            element.basix_element.entity_closure_dofs,  # type: ignore
            int(np.prod(element.value_shape)),  # type: ignore
        )
    else:
        cpp_dofmap = _cpp.fem.create_dofmap(
            mesh.comm,
            mesh.topology._cpp_object,
            element._cpp_object,  # type: ignore
        )
    assert np.issubdtype(mesh.geometry.x.dtype, element.dtype), (  # type: ignore
        "Mesh and element dtype are not compatible."
    )

    # Initialize the cpp.FunctionSpace
    try:
        cppV = _cpp.fem.FunctionSpace_float64(mesh._cpp_object, element._cpp_object, cpp_dofmap)  # type: ignore
    except TypeError:
        cppV = _cpp.fem.FunctionSpace_float32(mesh._cpp_object, element._cpp_object, cpp_dofmap)  # type: ignore

    return FunctionSpace(mesh, ufl_e, cppV)


class FunctionSpace(ufl.FunctionSpace, Generic[Real]):
    """A space on which Functions (fields) can be defined."""

    _cpp_object: _cpp.fem.FunctionSpace_float32 | _cpp.fem.FunctionSpace_float64
    _mesh: Mesh[Real]

    def __init__(
        self,
        mesh: Mesh[Real],
        element: ufl.finiteelement.AbstractFiniteElement,
        cppV: (_cpp.fem.FunctionSpace_float32 | _cpp.fem.FunctionSpace_float64),
    ):
        """Create a finite element function space.

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

        Args:
            mesh: Mesh that space is defined on.
            element: UFL finite element.
            cppV: Compiled C++ function space.
        """
        if mesh._cpp_object is not cppV.mesh:
            raise RuntimeError("Meshes do not match in function space initialisation.")
        ufl_domain = mesh.ufl_domain()
        self._cpp_object = cppV
        self._mesh = mesh
        super().__init__(ufl_domain, element)

    def clone(self) -> FunctionSpace[Real]:
        """Create a FunctionSpace which shares data with this space.

        The new space has a different unique integer ID.

        Create a new FunctionSpace :math:`W` which shares data with this
        FunctionSpace :math:`V`, but with a different unique integer ID.

        This function is helpful for defining mixed problems and using
        blocked linear algebra. For example, a matrix block defined on
        the spaces :math:`V \\times W` where, :math:`V` and :math:`W`
        are defined on the same finite element and mesh can be
        identified as an off-diagonal block whereas the :math:`V \\times
        V` and :math:`V \\times V` matrices can be identified as
        diagonal blocks. This is relevant for the handling of boundary
        conditions.

        Returns:
            A new function space that shares data
        """  # noqa: D301
        try:
            Vcpp = _cpp.fem.FunctionSpace_float64(
                self._cpp_object.mesh, self._cpp_object.element, self._cpp_object.dofmap
            )  # type: ignore
        except TypeError:
            Vcpp = _cpp.fem.FunctionSpace_float32(
                self._cpp_object.mesh, self._cpp_object.element, self._cpp_object.dofmap
            )  # type: ignore
        return FunctionSpace(self._mesh, self.ufl_element(), Vcpp)

    @property
    def num_sub_spaces(self) -> int:
        """Number of sub spaces."""
        return self.element.num_sub_elements

    def sub(self, i: int) -> FunctionSpace[Real]:
        """Return the i-th sub space.

        Args:
            i: Index of the subspace to extract.

        Returns:
            A subspace.

        Note:
            If the subspace is re-used, for performance reasons the
            returned subspace should be stored by the caller to avoid
            repeated re-computation of the subspace.
        """
        assert self.ufl_element().num_sub_elements > i
        sub_element = self.ufl_element().sub_elements[i]
        cppV_sub = self._cpp_object.sub([i])  # type: ignore
        return FunctionSpace(self._mesh, sub_element, cppV_sub)

    def component(self):
        """Return the component relative to the parent space."""
        return self._cpp_object.component()  # type: ignore

    def contains(self, V) -> bool:
        """Check if a space is contained in, or is the same as, this space.

        Args:
            V: The space to check to for inclusion.

        Returns:
           `` True`` if ``V`` is contained in, or is the same as, this
           space.
        """
        return self._cpp_object.contains(V._cpp_object)  # type: ignore

    def __eq__(self, other):
        """Comparison for equality."""
        return super().__eq__(other) and self._cpp_object == other._cpp_object

    def __ne__(self, other):
        """Comparison for inequality."""
        return super().__ne__(other) or self._cpp_object != other._cpp_object

    def ufl_function_space(self) -> ufl.FunctionSpace:
        """UFL function space."""
        return self

    @cached_property
    def element(self) -> FiniteElement[Real]:
        """Function space finite element."""
        return FiniteElement(self._cpp_object.element)

    @property
    def dofmap(self) -> DofMap:
        """Degree-of-freedom map associated with the function space."""
        return DofMap(self._cpp_object.dofmap)

    @property
    def dofmaps(self) -> list[DofMap]:
        """The geometry dofmaps, one per cell type."""
        return [DofMap(_o) for _o in self._cpp_object.dofmaps]

    @property
    def mesh(self) -> Mesh[Real]:
        """Mesh on which the function space is defined."""
        return self._mesh

    def collapse(self) -> tuple[FunctionSpace[Real], list[npt.NDArray[np.int32]]]:
        """Create a new function space by collapsing a subspace.

        Returns:
            A new function space and the map from new to old
            degrees-of-freedom.
        """
        cpp_space, dofs = self._cpp_object.collapse()
        V = FunctionSpace(self._mesh, self.ufl_element(), cpp_space)
        return V, dofs

    def tabulate_dof_coordinates(self) -> npt.NDArray[Real]:
        """Tabulate coordinates of function space degrees-of-freedom.

        Returns:
            Coordinates of the degrees-of-freedom.

        Note:
            This method is only for elements with point evaluation
            degrees-of-freedom.
        """
        return self._cpp_object.tabulate_dof_coordinates()
