# Copyright (C) 2017-2021 Chris N. Richardson, Garth N. Wells and
# Jørgen S. Dokken
#
# This file is part of DOLFINx (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later
"""Support for representing Dirichlet boundary conditions that are enforced
via modification of linear systems."""

from __future__ import annotations

import numbers
from collections.abc import Callable, Iterable

import numpy as np
import numpy.typing as npt

import dolfinx
from dolfinx import cpp as _cpp
from dolfinx.fem.function import Constant, Function, FunctionSpace


def locate_dofs_geometrical(
    V: dolfinx.fem.FunctionSpace | Iterable[dolfinx.fem.FunctionSpace],
    marker: Callable,
) -> np.ndarray:
    """Locate degrees-of-freedom geometrically using a marker function.

    Args:
        V: Function space(s) in which to search for degree-of-freedom
            indices.
        marker: A function that takes an array of points ``x`` with
            shape ``(gdim, num_points)`` and returns an array of
            booleans of length ``num_points``, evaluating to ``True``
            for entities whose degree-of-freedom should be returned.

    Returns:
        An array of degree-of-freedom indices (local to the process) for
        degrees-of-freedom whose coordinate evaluates to True for the
        marker function.

        If ``V`` is a list of two function spaces, then a 2-D array of
        shape (number of dofs, 2) is returned.

        Returned degree-of-freedom indices are unique and ordered by the
        first column.
    """
    if not isinstance(V, Iterable):
        return _cpp.fem.locate_dofs_geometrical(V._cpp_object, marker)  # type: ignore

    _V = [space._cpp_object for space in V]  # type: ignore
    return _cpp.fem.locate_dofs_geometrical(_V, marker)


def locate_dofs_topological(
    V: dolfinx.fem.FunctionSpace | Iterable[dolfinx.fem.FunctionSpace],
    entity_dim: int,
    entities: npt.NDArray[np.int32],
    remote: bool = True,
) -> np.ndarray:
    """Locate degrees-of-freedom belonging to mesh entities topologically.

    Args:
        V: Function space(s) in which to search for degree-of-freedom
            indices.
        entity_dim: Topological dimension of entities where
            degrees-of-freedom are located.
        entities: Indices of mesh entities of dimension ``entity_dim``
            where degrees-of-freedom are located.
        remote: True to return also "remotely located" degree-of-freedom
            indices.

    Returns:
        An array of degree-of-freedom indices (local to the process) for
        degrees-of-freedom topologically belonging to mesh entities.

        If ``V`` is a list of two function spaces, then a 2-D array of
        shape (number of dofs, 2) is returned.

        Returned degree-of-freedom indices are unique and ordered by the
        first column.
    """
    _entities = np.asarray(entities, dtype=np.int32)
    if not isinstance(V, Iterable):
        return _cpp.fem.locate_dofs_topological(V._cpp_object, entity_dim, _entities, remote)  # type: ignore

    _V = [space._cpp_object for space in V]  # type: ignore
    return _cpp.fem.locate_dofs_topological(_V, entity_dim, _entities, remote)


class DirichletBC:
    _cpp_object: (
        _cpp.fem.DirichletBC_complex64
        | _cpp.fem.DirichletBC_complex128
        | _cpp.fem.DirichletBC_float32
        | _cpp.fem.DirichletBC_float64
    )

    def __init__(self, bc):
        """Representation of Dirichlet boundary condition which is imposed
        on a linear system.

        Note:
            Dirichlet boundary conditions  should normally be
            constructed using :func:`fem.dirichletbc` and not using this
            class initialiser. This class is combined with different
            base classes that depend on the scalar type of the boundary
            condition.

        Args:
            value: Lifted boundary values function.
            dofs: Local indices of degrees of freedom in function space
                to which boundary condition applies. Expects array of
                size (number of dofs, 2) if function space of the
                problem, ``V`` is passed. Otherwise assumes function
                space of the problem is the same of function space of
                boundary values function.
            V: Function space of a problem to which boundary conditions
                are applied.
        """
        self._cpp_object = bc

    @property
    def g(self) -> Function | Constant | np.ndarray:
        """The boundary condition value(s)"""
        return self._cpp_object.value

    @property
    def function_space(self) -> dolfinx.fem.FunctionSpace:
        """The function space on which the boundary condition is defined"""
        return self._cpp_object.function_space

    def set(
        self, x: npt.NDArray, x0: npt.NDArray[np.int32] | None = None, alpha: float = 1
    ) -> None:
        """Set entries in an array that are constrained by Dirichlet
        boundary conditions.

        Entries in ``x`` that are constrained by a Dirichlet boundary
        conditions are set to ``alpha * (x_bc - x0)``, where ``x_bc`` is
        the (interpolated) boundary condition value.

        For elements with point-wise evaluated degrees-of-freedom, e.g.
        Lagrange elements, ``x_bc`` is the value of the boundary condition
        at the degree-of-freedom. For elements with moment
        degrees-of-freedom, ``x_bc`` is the value of the boundary condition
        interpolated into the finite element space.

        If `x` includes ghosted entries (entries available on the calling
        rank but owned by another rank), ghosted entries constrained by a
        Dirichlet condition will also be set.

        Args:
            x: Array to modify for Dirichlet boundary conditions.
            x0: Optional array used in computing the value to set. If
                not provided it is treated as zero.
            alpha: Scaling factor.
        """
        self._cpp_object.set(x, x0, alpha)

    def dof_indices(self) -> tuple[npt.NDArray[np.int32], int]:
        """Access dof indices `(local indices, unrolled)`, including
        ghosts, to which a Dirichlet condition is applied, and the index to
        the first non-owned (ghost) index. The array of indices is sorted.

        Note:
            The returned array is read-only.

        Returns:
            Sorted array of dof indices (unrolled) and index to the
            first entry in the dof index array that is not owned. Entries
            `dofs[:pos]` are owned and entries `dofs[pos:]` are ghosts.
        """
        return self._cpp_object.dof_indices()


def dirichletbc(
    value: Function | Constant | np.ndarray,
    dofs: npt.NDArray[np.int32],
    V: dolfinx.fem.FunctionSpace | None = None,
) -> DirichletBC:
    """Create a representation of Dirichlet boundary condition which
    is imposed on a linear system.

    Args:
        value: Lifted boundary values function. It must have a ``dtype``
            property.
        dofs: Local indices of degrees of freedom in function space to
            which boundary condition applies. Expects array of size
            (number of dofs, 2) if function space of the problem, ``V``,
            is passed. Otherwise assumes function space of the problem
            is the same of function space of boundary values function.
        V: Function space of a problem to which boundary conditions are
            applied.

    Returns:
        A representation of the boundary condition for modifying linear
        systems.
    """
    if isinstance(value, numbers.Number):
        value = np.asarray(value)

    try:
        dtype = value.dtype
        if np.issubdtype(dtype, np.float32):
            bctype = _cpp.fem.DirichletBC_float32
        elif np.issubdtype(dtype, np.float64):
            bctype = _cpp.fem.DirichletBC_float64
        elif np.issubdtype(dtype, np.complex64):
            bctype = _cpp.fem.DirichletBC_complex64
        elif np.issubdtype(dtype, np.complex128):
            bctype = _cpp.fem.DirichletBC_complex128
        else:
            raise NotImplementedError(f"Type {value.dtype} not supported.")
    except AttributeError:
        raise AttributeError("Boundary condition value must have a dtype attribute.")

    # Unwrap value object, if required
    if isinstance(value, np.ndarray):
        _value = value
    else:
        try:
            _value = value._cpp_object
        except AttributeError:
            _value = value  # type: ignore[assignment]

    if V is not None:
        try:
            bc = bctype(_value, dofs, V)
        except TypeError:
            bc = bctype(_value, dofs, V._cpp_object)
    else:
        bc = bctype(_value, dofs)

    return DirichletBC(bc)


def bcs_by_block(
    spaces: Iterable[FunctionSpace | None], bcs: Iterable[DirichletBC]
) -> list[list[DirichletBC]]:
    """Arrange Dirichlet boundary conditions by the function space that
    they constrain.

    Given a sequence of function spaces ``spaces`` and a sequence of
    DirichletBC objects ``bcs``, return a list where the ith entry is
    the list of DirichletBC objects whose space is contained in
    ``space[i]``.
    """

    def _bc_space(V, bcs):
        """Return list of bcs that have the same space as V"""
        return [bc for bc in bcs if V.contains(bc.function_space)]

    return [_bc_space(V, bcs) if V is not None else [] for V in spaces]