from typing import Optional
from enum import Enum

import warp.fem.domain as _domain
import warp.fem.geometry as _geometry
import warp.fem.polynomial as _polynomial

from .function_space import FunctionSpace
from .topology import SpaceTopology
from .basis_space import BasisSpace, PointBasisSpace
from .collocated_function_space import CollocatedFunctionSpace

from .grid_2d_function_space import (
    GridPiecewiseConstantBasis,
    GridBipolynomialBasisSpace,
    GridDGBipolynomialBasisSpace,
    GridSerendipityBasisSpace,
    GridDGSerendipityBasisSpace,
    GridDGPolynomialBasisSpace,
)
from .grid_3d_function_space import (
    GridTripolynomialBasisSpace,
    GridDGTripolynomialBasisSpace,
    Grid3DPiecewiseConstantBasis,
    Grid3DSerendipityBasisSpace,
    Grid3DDGSerendipityBasisSpace,
    Grid3DDGPolynomialBasisSpace,
)
from .trimesh_2d_function_space import (
    Trimesh2DPiecewiseConstantBasis,
    Trimesh2DPolynomialBasisSpace,
    Trimesh2DDGPolynomialBasisSpace,
    Trimesh2DNonConformingPolynomialBasisSpace,
)
from .tetmesh_function_space import (
    TetmeshPiecewiseConstantBasis,
    TetmeshPolynomialBasisSpace,
    TetmeshDGPolynomialBasisSpace,
    TetmeshNonConformingPolynomialBasisSpace,
)
from .quadmesh_2d_function_space import (
    Quadmesh2DPiecewiseConstantBasis,
    Quadmesh2DBipolynomialBasisSpace,
    Quadmesh2DDGBipolynomialBasisSpace,
    Quadmesh2DSerendipityBasisSpace,
    Quadmesh2DDGSerendipityBasisSpace,
    Quadmesh2DPolynomialBasisSpace,
)
from .hexmesh_function_space import (
    HexmeshPiecewiseConstantBasis,
    HexmeshTripolynomialBasisSpace,
    HexmeshDGTripolynomialBasisSpace,
    HexmeshSerendipityBasisSpace,
    HexmeshDGSerendipityBasisSpace,
    HexmeshPolynomialBasisSpace,
)

from .partition import SpacePartition, make_space_partition
from .restriction import SpaceRestriction


from .dof_mapper import DofMapper, IdentityMapper, SymmetricTensorMapper, SkewSymmetricTensorMapper


def make_space_restriction(
    space: Optional[FunctionSpace] = None,
    space_partition: Optional[SpacePartition] = None,
    domain: Optional[_domain.GeometryDomain] = None,
    space_topology: Optional[SpaceTopology] = None,
    device=None,
    temporary_store: "Optional[warp.fem.cache.TemporaryStore]" = None,
) -> SpaceRestriction:
    """
    Restricts a function space partition to a Domain, i.e. a subset of its elements.

    One of `space_partition`, `space_topology`, or `space` must be provided (and will be considered in that order).

    Args:
        space: (deprecated) if neither `space_partition` nor `space_topology` are provided, the space defining the topology to restrict
        space_partition: the subset of nodes from the space topology to consider
        domain: the domain to restrict the space to, defaults to all cells of the space geometry or partition.
        space_topology: the space topology to be restricted, if `space_partition` is ``None``.
        device: device on which to perform and store computations
        temporary_store: shared pool from which to allocate temporary arrays
    """

    if space_partition is None:
        if space_topology is None:
            assert space is not None
            space_topology = space.topology

        if domain is None:
            domain = _domain.Cells(geometry=space_topology.geometry)

        space_partition = make_space_partition(
            space_topology=space_topology, geometry_partition=domain.geometry_partition
        )
    elif domain is None:
        domain = _domain.Cells(geometry=space_partition.geo_partition)

    return SpaceRestriction(
        space_partition=space_partition, domain=domain, device=device, temporary_store=temporary_store
    )


class ElementBasis(Enum):
    """Choice of basis function to equip individual elements"""

    LAGRANGE = 0
    """Lagrange basis functions :math:`P_k` for simplices, tensor products :math:`Q_k` for squares and cubes"""
    SERENDIPITY = 1
    """Serendipity elements :math:`S_k`, corresponding to Lagrange nodes with interior points removed (for degree <= 3)"""
    NONCONFORMING_POLYNOMIAL = 2
    """Simplex Lagrange basis functions :math:`P_{kd}` embedded into non conforming reference elements (e.g. squares or cubes). Discontinuous only."""


def make_polynomial_basis_space(
    geo: _geometry.Geometry,
    degree: int = 1,
    element_basis: Optional[ElementBasis] = None,
    discontinuous: bool = False,
    family: Optional[_polynomial.Polynomial] = None,
) -> BasisSpace:
    """
    Equips a geometry with a polynomial basis.

    Args:
        geo: the Geometry on which to build the space
        degree: polynomial degree of the per-element shape functions
        discontinuous: if True, use Discontinuous Galerkin shape functions. Discontinuous is implied if degree is 0, i.e, piecewise-constant shape functions.
        element_basis: type of basis function for the individual elements
        family: Polynomial family used to generate the shape function basis. If not provided, a reasonable basis is chosen.

    Returns:
        the constructed basis space
    """

    base_geo = geo.base if isinstance(geo, _geometry.DeformedGeometry) else geo

    if element_basis is None:
        element_basis = ElementBasis.LAGRANGE

    if isinstance(base_geo, _geometry.Grid2D):
        if degree == 0:
            return GridPiecewiseConstantBasis(geo)

        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
            if discontinuous:
                return GridDGSerendipityBasisSpace(geo, degree=degree, family=family)
            else:
                return GridSerendipityBasisSpace(geo, degree=degree, family=family)

        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
            return GridDGPolynomialBasisSpace(geo, degree=degree)

        if discontinuous:
            return GridDGBipolynomialBasisSpace(geo, degree=degree, family=family)
        else:
            return GridBipolynomialBasisSpace(geo, degree=degree, family=family)

    if isinstance(base_geo, _geometry.Grid3D):
        if degree == 0:
            return Grid3DPiecewiseConstantBasis(geo)

        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
            if discontinuous:
                return Grid3DDGSerendipityBasisSpace(geo, degree=degree, family=family)
            else:
                return Grid3DSerendipityBasisSpace(geo, degree=degree, family=family)

        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
            return Grid3DDGPolynomialBasisSpace(geo, degree=degree)

        if discontinuous:
            return GridDGTripolynomialBasisSpace(geo, degree=degree, family=family)
        else:
            return GridTripolynomialBasisSpace(geo, degree=degree, family=family)

    if isinstance(base_geo, _geometry.Trimesh2D):
        if degree == 0:
            return Trimesh2DPiecewiseConstantBasis(geo)

        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
            raise NotImplementedError("Serendipity variant not implemented yet")

        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
            return Trimesh2DNonConformingPolynomialBasisSpace(geo, degree=degree)

        if discontinuous:
            return Trimesh2DDGPolynomialBasisSpace(geo, degree=degree)
        else:
            return Trimesh2DPolynomialBasisSpace(geo, degree=degree)

    if isinstance(base_geo, _geometry.Tetmesh):
        if degree == 0:
            return TetmeshPiecewiseConstantBasis(geo)

        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
            raise NotImplementedError("Serendipity variant not implemented yet")

        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
            return TetmeshNonConformingPolynomialBasisSpace(geo, degree=degree)

        if discontinuous:
            return TetmeshDGPolynomialBasisSpace(geo, degree=degree)
        else:
            return TetmeshPolynomialBasisSpace(geo, degree=degree)

    if isinstance(base_geo, _geometry.Quadmesh2D):
        if degree == 0:
            return Quadmesh2DPiecewiseConstantBasis(geo)

        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
            if discontinuous:
                return Quadmesh2DDGSerendipityBasisSpace(geo, degree=degree, family=family)
            else:
                return Quadmesh2DSerendipityBasisSpace(geo, degree=degree, family=family)

        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
            return Quadmesh2DPolynomialBasisSpace(geo, degree=degree)

        if discontinuous:
            return Quadmesh2DDGBipolynomialBasisSpace(geo, degree=degree, family=family)
        else:
            return Quadmesh2DBipolynomialBasisSpace(geo, degree=degree, family=family)

    if isinstance(base_geo, _geometry.Hexmesh):
        if degree == 0:
            return HexmeshPiecewiseConstantBasis(geo)

        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
            if discontinuous:
                return HexmeshDGSerendipityBasisSpace(geo, degree=degree, family=family)
            else:
                return HexmeshSerendipityBasisSpace(geo, degree=degree, family=family)

        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
            return HexmeshPolynomialBasisSpace(geo, degree=degree)

        if discontinuous:
            return HexmeshDGTripolynomialBasisSpace(geo, degree=degree, family=family)
        else:
            return HexmeshTripolynomialBasisSpace(geo, degree=degree, family=family)

    raise NotImplementedError()


def make_collocated_function_space(
    basis_space: BasisSpace, dtype: type = float, dof_mapper: Optional[DofMapper] = None
) -> CollocatedFunctionSpace:
    """
    Constructs a function space from a basis space and a value type, such that all degrees of freedom of the value type are stored at each of the basis nodes.

    Args:
        geo: the Geometry on which to build the space
        dtype: value type the function space. If ``dof_mapper`` is provided, the value type from the DofMapper will be used instead.
        dof_mapper: mapping from node degrees of freedom to function values, defaults to Identity. Useful for reduced coordinates, e.g. :py:class:`SymmetricTensorMapper` maps 2x2 (resp 3x3) symmetric tensors to 3 (resp 6) degrees of freedom.

    Returns:
        the constructed function space
    """
    return CollocatedFunctionSpace(basis_space, dtype=dtype, dof_mapper=dof_mapper)


def make_polynomial_space(
    geo: _geometry.Geometry,
    dtype: type = float,
    dof_mapper: Optional[DofMapper] = None,
    degree: int = 1,
    element_basis: Optional[ElementBasis] = None,
    discontinuous: bool = False,
    family: Optional[_polynomial.Polynomial] = None,
) -> CollocatedFunctionSpace:
    """
    Equips a geometry with a collocated, polynomial function space.
    Equivalent to successive calls to :func:`make_polynomial_basis_space` and `make_collocated_function_space`.

    Args:
        geo: the Geometry on which to build the space
        dtype: value type the function space. If ``dof_mapper`` is provided, the value type from the DofMapper will be used instead.
        dof_mapper: mapping from node degrees of freedom to function values, defaults to Identity. Useful for reduced coordinates, e.g. :py:class:`SymmetricTensorMapper` maps 2x2 (resp 3x3) symmetric tensors to 3 (resp 6) degrees of freedom.
        degree: polynomial degree of the per-element shape functions
        discontinuous: if True, use Discontinuous Galerkin shape functions. Discontinuous is implied if degree is 0, i.e, piecewise-constant shape functions.
        element_basis: type of basis function for the individual elements
        family: Polynomial family used to generate the shape function basis. If not provided, a reasonable basis is chosen.

    Returns:
        the constructed function space
    """

    basis_space = make_polynomial_basis_space(geo, degree, element_basis, discontinuous, family)
    return CollocatedFunctionSpace(basis_space, dtype=dtype, dof_mapper=dof_mapper)