# Source code for ufl.domain

    """Types for representing a geometric domain."""

    # Copyright (C) 2008-2016 Martin Sandve Alnæs
    #
    # This file is part of UFL (https://www.fenicsproject.org)
    #
    # SPDX-License-Identifier:    LGPL-3.0-or-later

    from __future__ import annotations  # To avoid cyclic import when type-hinting.

    import numbers
    from collections.abc import Iterable, Sequence
    from typing import TYPE_CHECKING

    if TYPE_CHECKING:
        from ufl.core.expr import Expr
        from ufl.finiteelement import AbstractFiniteElement  # To avoid cyclic import when type-hinting.
        from ufl.form import Form
    from ufl.cell import AbstractCell
    from ufl.core.ufl_id import attach_ufl_id
    from ufl.core.ufl_type import UFLObject
    from ufl.corealg.traversal import traverse_unique_terminals
    from ufl.sobolevspace import H1

    # Export list for ufl.classes
    __all_classes__ = ["AbstractDomain", "Mesh", "MeshView"]

    class AbstractDomain:
        """Symbolic representation of a geometric domain.

        Domain has only a geometric and a topological dimension.
        """

        def __init__(self, topological_dimension, geometric_dimension):
            """Initialise."""
            # Validate dimensions
            if not isinstance(geometric_dimension, numbers.Integral):
                raise ValueError(
                    f"Expecting integer geometric dimension, not {geometric_dimension.__class__}"
                )
            if not isinstance(topological_dimension, numbers.Integral):
                raise ValueError(
                    f"Expecting integer topological dimension, not {topological_dimension.__class__}"
                )
            if topological_dimension > geometric_dimension:
                raise ValueError("Topological dimension cannot be larger than geometric dimension.")

            # Store validated dimensions
            self._topological_dimension = topological_dimension
            self._geometric_dimension = geometric_dimension

        def geometric_dimension(self):
            """Return the dimension of the space this domain is embedded in."""
            return self._geometric_dimension

        def topological_dimension(self):
            """Return the dimension of the topology of this domain."""
            return self._topological_dimension

        @property
        def meshes(self):
            """Return the component meshes."""
            raise NotImplementedError("meshes() method not implemented")

        def __len__(self):
            """Return number of component meshes."""
            return len(self.meshes)

        def __getitem__(self, i):
            """Return i-th component mesh."""
            if i >= len(self):
                raise ValueError(f"index ({i}) >= num. component meshes ({len(self)})")
            return self.meshes[i]

        def __iter__(self):
            """Return iterable component meshes."""
            return iter(self.meshes)

        def iterable_like(self, element: AbstractFiniteElement) -> Iterable[Mesh] | MeshSequence:
            """Return iterable object that is iterable like ``element``."""
            raise NotImplementedError("iterable_like() method not implemented")

        def can_make_function_space(self, element: AbstractFiniteElement) -> bool:
            """Check whether this mesh can make a function space with ``element``."""
            raise NotImplementedError("can_make_function_space() method not implemented")

    # TODO: Would it be useful to have a domain representing R^d? E.g. for
    # Expression.
    # class EuclideanSpace(AbstractDomain):
    #     def __init__(self, geometric_dimension):
    #         AbstractDomain.__init__(self, geometric_dimension, geometric_dimension)

    @attach_ufl_id
    class Mesh(AbstractDomain, UFLObject):
        """Symbolic representation of a mesh."""

        def __init__(self, coordinate_element, ufl_id=None, cargo=None):
            """Initialise."""
            self._ufl_id = self._init_ufl_id(ufl_id)

            # Store reference to object that will not be used by UFL
            self._ufl_cargo = cargo
            if cargo is not None and cargo.ufl_id() != self._ufl_id:
                raise ValueError("Expecting cargo object (e.g. dolfin.Mesh) to have the same ufl_id.")

            # No longer accepting coordinates provided as a Coefficient
            from ufl.coefficient import Coefficient

            if isinstance(coordinate_element, Coefficient | AbstractCell):
                raise ValueError("Expecting a coordinate element in the ufl.Mesh construct.")

            # Store coordinate element
            self._ufl_coordinate_element = coordinate_element

            # Derive dimensions from element
            (gdim,) = coordinate_element.reference_value_shape
            tdim = coordinate_element.cell.topological_dimension()
            AbstractDomain.__init__(self, tdim, gdim)

        def ufl_cargo(self):
            """Return carried object that will not be used by UFL."""
            return self._ufl_cargo

        def ufl_coordinate_element(self):
            """Get the coordinate element."""
            return self._ufl_coordinate_element

        def ufl_cell(self):
            """Get the cell."""
            return self._ufl_coordinate_element.cell

        def is_piecewise_linear_simplex_domain(self):
            """Check if the domain is a piecewise linear simplex."""
            ce = self._ufl_coordinate_element
            return ce.embedded_superdegree <= 1 and ce in H1 and self.ufl_cell().is_simplex()

        def __repr__(self):
            """Representation."""
            r = f"Mesh({self._ufl_coordinate_element!r}, {self._ufl_id!r})"
            return r

        def __str__(self):
            """Format as a string."""
            return f"<Mesh #{self._ufl_id}>"

        def _ufl_hash_data_(self):
            """UFL hash data."""
            return (self._ufl_id, self._ufl_coordinate_element)

        def _ufl_signature_data_(self, renumbering):
            """UFL signature data."""
            return ("Mesh", renumbering[self], self._ufl_coordinate_element)

        # NB! Dropped __lt__ here, don't want users to write 'mesh1 <
        # mesh2'.
        def _ufl_sort_key_(self):
            """UFL sort key."""
            typespecific = (self._ufl_id, self._ufl_coordinate_element)
            return (self.geometric_dimension(), self.topological_dimension(), "Mesh", typespecific)

        @property
        def meshes(self):
            """Return the component meshes."""
            return (self,)

        def iterable_like(self, element: AbstractFiniteElement) -> Iterable[Mesh]:
            """Return iterable object that is iterable like ``element``."""
            return iter(self for _ in range(element.num_sub_elements))

        def can_make_function_space(self, element: AbstractFiniteElement) -> bool:
            """Check whether this mesh can make a function space with ``element``."""
            # Can use with any element.
            return True

    class MeshSequence(AbstractDomain, UFLObject):
        """Symbolic representation of a mixed mesh.

        This class represents a collection of meshes that, along with
        a :class:`MixedElement`, represent a mixed function space defined on
        multiple domains. This abstraction allows for defining the
        mixed function space with the conventional :class:`FunctionSpace`
        class and integrating multi-domain problems seamlessly.

        Currently, all component meshes must have the same cell type (and
        thus the same topological dimension).

        Currently, one can only perform cell integrations when
        :class:`MeshSequence` is used.

        .. code-block:: python3

            cell = triangle
            mesh0 = Mesh(FiniteElement("Lagrange", cell, 1, (2,), identity_pullback, H1))
            mesh1 = Mesh(FiniteElement("Lagrange", cell, 1, (2,), identity_pullback, H1))
            domain = MeshSequence([mesh0, mesh1])
            elem0 = FiniteElement("Lagrange", cell, 1, (), identity_pullback, H1)
            elem1 = FiniteElement("Lagrange", cell, 2, (), identity_pullback, H1)
            elem = MixedElement([elem0, elem1])
            V = FunctionSpace(domain, elem)
            v = TestFunction(V)
            v0, v1 = split(v)

        """

        def __init__(self, meshes: Sequence[Mesh]):
            """Initialise."""
            if any(isinstance(m, MeshSequence) for m in meshes):
                raise NotImplementedError("""
                    Currently component meshes can not include MeshSequence instances""")
            # currently only support single cell type.
            (self._ufl_cell,) = set(m.ufl_cell() for m in meshes)
            (gdim,) = set(m.geometric_dimension() for m in meshes)
            # TODO: Need to change for more general mixed meshes.
            (tdim,) = set(m.topological_dimension() for m in meshes)
            AbstractDomain.__init__(self, tdim, gdim)
            self._meshes = tuple(meshes)

        def ufl_cell(self):
            """Get the cell."""
            # TODO: Might need MixedCell class for more general mixed meshes.
            return self._ufl_cell

        def __repr__(self):
            """Representation."""
            return f"MeshSequence({self._meshes!r})"

        def __str__(self):
            """Format as a string."""
            return f"<MeshSequence #{self._meshes}>"

        def _ufl_hash_data_(self):
            """UFL hash data."""
            return ("MeshSequence", tuple(m._ufl_hash_data_() for m in self._meshes))

        def _ufl_signature_data_(self, renumbering):
            """UFL signature data."""
            return ("MeshSequence", tuple(m._ufl_signature_data_(renumbering) for m in self._meshes))

        def _ufl_sort_key_(self):
            """UFL sort key."""
            return ("MeshSequence", tuple(m._ufl_sort_key_() for m in self._meshes))

        @property
        def meshes(self):
            """Return the component meshes."""
            return self._meshes

        def iterable_like(self, element: AbstractFiniteElement) -> MeshSequence:
            """Return iterable object that is iterable like ``element``."""
            if len(self) != element.num_sub_elements:
                raise RuntimeError(f"""len(self) ({len(self)}) !=
                    element.num_sub_elements ({element.num_sub_elements})""")
            return self

        def can_make_function_space(self, element: AbstractFiniteElement) -> bool:
            """Check whether this mesh can make a function space with ``element``."""
            if len(self) != element.num_sub_elements:
                return False
            else:
                return all(d.can_make_function_space(e) for d, e in zip(self, element.sub_elements))

    @attach_ufl_id
    class MeshView(AbstractDomain, UFLObject):
        """Symbolic representation of a mesh."""

        def __init__(self, mesh, topological_dimension, ufl_id=None):
            """Initialise."""
            self._ufl_id = self._init_ufl_id(ufl_id)

            # Store mesh
            self._ufl_mesh = mesh

            # Derive dimensions from element
            coordinate_element = mesh.ufl_coordinate_element()
            (gdim,) = coordinate_element.value_shape
            tdim = coordinate_element.cell.topological_dimension()
            AbstractDomain.__init__(self, tdim, gdim)

        def ufl_mesh(self):
            """Get the mesh."""
            return self._ufl_mesh

        def ufl_cell(self):
            """Get the cell."""
            return self._ufl_mesh.ufl_cell()

        def is_piecewise_linear_simplex_domain(self):
            """Check if the domain is a piecewise linear simplex."""
            return self._ufl_mesh.is_piecewise_linear_simplex_domain()

        def __repr__(self):
            """Representation."""
            tdim = self.topological_dimension()
            r = f"MeshView({self._ufl_mesh!r}, {tdim!r}, {self._ufl_id!r})"
            return r

        def __str__(self):
            """Format as a string."""
            return (
                f"<MeshView #{self._ufl_id} of dimension {self.topological_dimension()} over"
                f" mesh {self._ufl_mesh}>"
            )

        def _ufl_hash_data_(self):
            """UFL hash data."""
            return (self._ufl_id,) + self._ufl_mesh._ufl_hash_data_()

        def _ufl_signature_data_(self, renumbering):
            """UFL signature data."""
            return ("MeshView", renumbering[self], self._ufl_mesh._ufl_signature_data_(renumbering))

        # NB! Dropped __lt__ here, don't want users to write 'mesh1 <
        # mesh2'.
        def _ufl_sort_key_(self):
            """UFL sort key."""
            typespecific = (self._ufl_id, self._ufl_mesh)
            return (self.geometric_dimension(), self.topological_dimension(), "MeshView", typespecific)

    def as_domain(domain):
        """Convert any valid object to an AbstractDomain type."""
        if isinstance(domain, AbstractDomain):
            # Modern UFL files and dolfin behaviour
            (domain,) = set(domain.meshes)
            return domain
        try:
            return extract_unique_domain(domain)
        except AttributeError:
            domain = domain.ufl_domain()
            (domain,) = set(domain.meshes)
            return domain

    def sort_domains(domains: Sequence[AbstractDomain]):
        """Sort domains in a canonical ordering.

        Args:
            domains: Sequence of domains.

        Returns:
            `tuple` of sorted domains.

        """
        return tuple(sorted(domains, key=lambda domain: domain._ufl_sort_key_()))

    def join_domains(domains: Sequence[AbstractDomain], expand_mesh_sequence: bool = True):
        """Take a list of domains and return a set with only unique domain objects.

        Args:
            domains: Sequence of domains.
            expand_mesh_sequence: If True, MeshSequence components are expanded.

        Returns:
            `set` of domains.

        """
        # Use hashing to join domains, ignore None
        joined_domains = set(domains) - set((None,))
        if expand_mesh_sequence:
            unrolled_joined_domains = set()
            for domain in joined_domains:
                unrolled_joined_domains.update(domain.meshes)
            joined_domains = unrolled_joined_domains

        if not joined_domains:
            return ()

        # Check geometric dimension compatibility
        gdims = set()
        for domain in joined_domains:
            gdims.add(domain.geometric_dimension())
        if len(gdims) != 1:
            raise ValueError("Found domains with different geometric dimensions.")

        return joined_domains

    # TODO: Move these to an analysis module?

    def extract_domains(expr: Expr | Form, expand_mesh_sequence: bool = True):
        """Return all domains expression is defined on.

        Args:
            expr: Expr or Form.
            expand_mesh_sequence: If True, MeshSequence components are expanded.

        Returns:
            `tuple` of domains.

        """
        from ufl.form import Form

        if isinstance(expr, Form):
            if not expand_mesh_sequence:
                raise NotImplementedError("""
                    Currently, can only extract domains from a Form with expand_mesh_sequence=True""")
            # Be consistent with the numbering used in signature.
            return tuple(expr.domain_numbering().keys())
        else:
            domainlist = []
            for t in traverse_unique_terminals(expr):
                domainlist.extend(t.ufl_domains())
            return sort_domains(join_domains(domainlist, expand_mesh_sequence=expand_mesh_sequence))

    def extract_unique_domain(expr, expand_mesh_sequence: bool = True):
        """Return the single unique domain expression is defined on or throw an error.

        Args:
            expr: Expr or Form.
            expand_mesh_sequence: If True, MeshSequence components are expanded.

        Returns:
            domain.

        """
        domains = extract_domains(expr, expand_mesh_sequence=expand_mesh_sequence)
        if len(domains) == 1:
            return domains[0]
        elif domains:
            raise ValueError("Found multiple domains, cannot return just one.")
        else:
            return None

    def find_geometric_dimension(expr):
        """Find the geometric dimension of an expression."""
        gdims = set()
        for t in traverse_unique_terminals(expr):
            # Can have multiple domains of the same cell type.
            domains = extract_domains(t)
            if len(domains) > 0:
                (gdim,) = set(domain.geometric_dimension() for domain in domains)
                gdims.add(gdim)

        if len(gdims) != 1:
            raise ValueError("Cannot determine geometric dimension from expression.")
        (gdim,) = gdims
        return gdim