from collections import abc
from abc import ABCMeta


class Node(abc.Sequence):
    """Node of a mesh.

    A node is a single point in 2D or 3D space. It is associated with a tag. The node is
    defined either as a point in 2D space with the `x` and `y` coordinates or in 3D
    space with `x`, `y`, and `z` coordinates. A node is the most basic constituent of a
    mesh. Nodes behave similar to tuples of length two or three: they are iterable,
    immutable, and hashable. A comparison also compares tags.
    """

    def __init__(self, coordinate, *, tag):
        """Initialize node.

        Args:
            coordinate (Iterable): coordinates defining the x, y(, and z) component
            tag (int): identifier
        """
        self._tag = int(tag)
        self._x, self._y, *z = map(float, coordinate)
        self._sdim = 2 + len(z)
        if self._sdim == 2:
            self._z = 0.0
        elif self._sdim == 3:
            self._z = z[0]
        else:
            raise ValueError(f"space dimension is '{self._sdim}' but expected 2 or 3")

    @property
    def x(self):
        """float: X coordinate."""
        return self._x

    @property
    def y(self):
        """float: Y coordinate."""
        return self._y

    @property
    def z(self):
        """float: Z coordinate."""
        return self._z

    @property
    def tag(self):
        """int: Identifier."""
        return self._tag

    @property
    def sdim(self):
        """int: Spatial dimension."""
        return self._sdim

    @property
    def in_3d(self):
        """tuple: Coordinates embedded in 3D space."""
        return self._x, self._y, self._z

    def __repr__(self):
        return f"Node({self[:]}, tag={self.tag})"

    def __len__(self):
        return self._sdim

    def __getitem__(self, key):
        """ """
        if self.sdim == 2:
            return (self.x, self.y)[key]
        return (self.x, self.y, self.z)[key]

    def __hash__(self):
        return hash(self[:] + (self.tag,))

    def __eq__(self, other):
        return self.tag == other.tag and self[:] == other[:]


class _ElementMeta(ABCMeta):
    """Metaclass for :class:`Element`.

    Require the attributes `dim` and `n_nodes`. If not specified, derive the COMSOL and
    JCMsuite names from the class name. The metaclass has to inherit from `ABCMeta` to
    keep the abstract base class `Sequence` working.
    """

    def __new__(cls, clsname, bases, attrs, /, **kwargs):
        for attr in ("dim", "n_nodes"):
            if attr not in attrs:
                TypeError(f"can't create class {clsname} without attribute {attr}.")
        dim = attrs["dim"]
        attrs.setdefault("comsol_name", "3 " + clsname[:3].lower())
        jcm_name = attrs.setdefault("jcm_name", clsname)
        attrs.setdefault("jcm_short_name", "#" + jcm_name[:2])
        attrs["jcm_info"] = {
            dim: f"<I>N{jcm_name}s",
            dim + 1: f"<I>NBoundary{jcm_name}s",
        }
        pre_bulk = f"# {jcm_name}s ###################\n# {jcm_name}s (p"
        pre_boundary = f"# {jcm_name} (p"
        middle = ", p".join(map(str, range(1, attrs["n_nodes"] + 1)))
        post_bulk = ", id)\n"
        post_boundary = (
            ", boundary id, topology indicator, (identic points))\n"  # codespell:ignore
        )
        attrs["jcm_header"] = {
            dim: pre_bulk + middle + post_bulk,
            dim + 1: pre_boundary + middle + post_boundary,
        }
        return super().__new__(cls, clsname, bases, attrs, **kwargs)


class Element(abc.Sequence, metaclass=_ElementMeta):
    """Mesh element.

    Elements are a collection of :class:`Node` that define a specific shape. The element
    itself defines the shape and it is assumed that the nodes are set accordingly. The
    dimension of the element is the dimension of the shape not the spatial dimension.
    Elements are iterable, immutable, and hashable. A comparison also compares tags.
    each class should specify attributes that describe how they are represented in
    different mesh files.
    """

    dim = -1
    n_nodes = 0

    def __init__(self, nodes, *, tag):
        """Initialize element.

        Args:
            nodes (Iterable): nodes of the element
            tag (int): identifier
        """
        self._sdim = nodes[0].sdim
        if self.sdim < self.dim:
            raise ValueError("node dimension must be larger than element dimension")
        for node in nodes:
            if node.sdim != self.sdim:
                raise ValueError("element contains nodes of mixed dimension")
        self._nodes = tuple(nodes)
        self._tag = int(tag)
        self.jcm_name = type(self).__name__

    @property
    def tag(self):
        """int: Identifier."""
        return self._tag

    @property
    def sdim(self):
        """int: Spatial dimension."""
        return self._sdim

    def __repr__(self):
        return f"{type(self).__name__}({self._nodes}, tag={self.tag})"

    def __len__(self):
        return self.n_nodes

    def __getitem__(self, key):
        return self._nodes[key]

    def __hash__(self):
        return hash((self.__class__, self.tag) + self._nodes)

    def __eq__(self, other):
        return (
            self.__class__ == other.__class__
            and self.tag == other.tag
            and all(a == b for a, b in zip(self, other))
        )


class Vertex(Element):
    """Actually just a node.

    The vertex gives another name to a node. As such it is a 0D mesh element. It exists
    in COMSOL meshes. In gmsh meshes it exists somewhat implicitly as `Point`.
    """

    dim = 0
    n_nodes = 1
    # It gets an gmsh type of None because it exists in gmsh only as `Point`, which is
    # an Entity (i.e. a Domain) but not as Element
    gmsh_type = None
    comsol_name = "3 vtx"


class Edge(Element):
    """Two nodes connected.

    An edge is a 1D mesh element. It exists in JCMsuite only as a boundary element of
    2D meshes.
    """

    dim = 1
    n_nodes = 2
    gmsh_type = 1
    jcm_short_name = "#E"


class Triangle(Element):
    """Three nodes that are not on a straight line.

    Triangles are the central element for 2D meshes.
    """

    dim = 2
    n_nodes = 3
    gmsh_type = 2


class Quadrilateral(Element):
    """Four nodes in a plane.

    The order of the nodes is different in different programs. In JCMsuite and gmsh they
    are in clockwise order, in COMSOL the last two nodes are exchanged.
    """

    dim = 2
    n_nodes = 4
    gmsh_type = 3
    comsol_name = "4 quad"
    jcm_short_name = "#Q"


class Tetrahedron(Element):
    """Four nodes that are not in a plane."""

    dim = 3
    n_nodes = 4
    gmsh_type = 4


class Pyramid(Element):
    """Five nodes forming a quadrilateral and four triangles.

    Like for the :class:`Quadrilateral`, COMSOL uses a different node order for the
    base. The last node is always the pyramid top.
    """

    dim = 3
    n_nodes = 5
    gmsh_type = 7


class Prism(Element):
    """Six nodes forming two triangles separated by some distance."""

    dim = 3
    n_nodes = 6
    gmsh_type = 6
    comsol_name = "5 prism"


class Hexahedron(Element):
    """Eight nodes forming two quadrilaterals separated by some distance.

    Like for the :class:`Quadrilateral`, COMSOL defines them with a different node
    order. JCMsuite, inconsistently, does so as well.
    """

    dim = 3
    n_nodes = 8
    gmsh_type = 5
    jcm_name = "Brick"
    jcm_short_name = "#B"


class Domain(abc.Set):
    """Domain of a mesh.

    A domain is the highest constituent of a mesh. It collects multiple elements of the
    same dimension that form a common shape or physical object together. It is
    essentially an immutable (and, thus, hashable) set. The order of the elements
    doesn't matter for the set operations, however, the elements are sorted by their tag
    upon initialization by their tag, to provide a fixed order when writing the mesh to a
    file. A domain can be empty, then its dimension and spatial dimension have to be
    given. In gmsh domains are named `Entity`.
    """

    def __init__(self, elements, *, tag, dim=None, sdim=None):
        """Initialize domain.

        Args:
            elements (Iterable): constituents of the domain
            tag (int): identifier
            dim (int, optional): dimension
            sdim (int, optional): spation dimension
        """
        if len(elements) != len(set(elements)):
            raise ValueError("duplicate elements")
        self._elements = sorted(elements, key=lambda item: item.tag)
        dims = {e.dim for e in self}
        if dim is not None:
            dims.add(dim)
        sdims = {e.sdim for e in self}
        if sdim is not None:
            sdims.add(sdim)
        if len(dims) != 1 or len(sdims) != 1:
            raise ValueError("invalid dimensions")
        self._dim = next(iter(dims))
        self._sdim = next(iter(sdims))
        self._tag = int(tag)

    @property
    def tag(self):
        """int: Identifier."""
        return self._tag

    @property
    def dim(self):
        """int: Dimension of the domain."""
        return self._dim

    @property
    def sdim(self):
        """int: Spatial dimension."""
        return self._sdim

    def __repr__(self):
        return "Domain({...}, " f"tag={self.tag}, dim={self.dim}, sdim={self.sdim})"

    def __iter__(self):
        return iter(self._elements)

    def __len__(self):
        return len(self._elements)

    def __contains__(self, item):
        return item in self._elements

    def __hash__(self):
        return hash((self._hash(), self.tag))

    def __le__(self, other):
        return self.tag == other.tag and super().__le__(other)

    def __ge__(self, other):
        return self.tag == other.tag and super().__le__(other)
