# Copyright (C) 2020-2023 Jørgen S. Dokken
#
# This file is part of DOLFINX_MPC
#
# SPDX-License-Identifier:    MIT
from __future__ import annotations

from typing import Callable, Dict, List, Optional, Tuple, Union

from petsc4py import PETSc as _PETSc

import dolfinx.cpp as _cpp
import dolfinx.fem as _fem
import dolfinx.mesh as _mesh
import numpy
import numpy.typing as npt
from dolfinx import default_scalar_type

import dolfinx_mpc.cpp

from .dictcondition import create_dictionary_constraint

_mpc_classes = Union[
    dolfinx_mpc.cpp.mpc.MultiPointConstraint_double,
    dolfinx_mpc.cpp.mpc.MultiPointConstraint_float,
    dolfinx_mpc.cpp.mpc.MultiPointConstraint_complex_double,
    dolfinx_mpc.cpp.mpc.MultiPointConstraint_complex_float,
]
_float_classes = Union[numpy.float32, numpy.float64, numpy.complex128, numpy.complex64]
_float_array_types = Union[
    npt.NDArray[numpy.float32],
    npt.NDArray[numpy.float64],
    npt.NDArray[numpy.complex64],
    npt.NDArray[numpy.complex128],
]
_mpc_data_classes = Union[
    dolfinx_mpc.cpp.mpc.mpc_data_double,
    dolfinx_mpc.cpp.mpc.mpc_data_float,
    dolfinx_mpc.cpp.mpc.mpc_data_complex_double,
    dolfinx_mpc.cpp.mpc.mpc_data_complex_float,
]


class MPCData:
    _cpp_object: _mpc_data_classes

    def __init__(
        self,
        slaves: npt.NDArray[numpy.int32],
        masters: npt.NDArray[numpy.int64],
        coeffs: _float_array_types,
        owners: npt.NDArray[numpy.int32],
        offsets: npt.NDArray[numpy.int32],
    ):
        if coeffs.dtype.type == numpy.float32:
            self._cpp_object = dolfinx_mpc.cpp.mpc.mpc_data_float(slaves, masters, coeffs, owners, offsets)
        elif coeffs.dtype.type == numpy.float64:
            self._cpp_object = dolfinx_mpc.cpp.mpc.mpc_data_double(slaves, masters, coeffs, owners, offsets)
        elif coeffs.dtype.type == numpy.complex64:
            self._cpp_object = dolfinx_mpc.cpp.mpc.mpc_data_complex_float(slaves, masters, coeffs, owners, offsets)
        elif coeffs.dtype.type == numpy.complex128:
            self._cpp_object = dolfinx_mpc.cpp.mpc.mpc_data_complex_double(slaves, masters, coeffs, owners, offsets)
        else:
            raise ValueError("Unsupported dtype {coeffs.dtype.type} for coefficients")

    @property
    def slaves(self):
        return self._cpp_object.slaves

    @property
    def masters(self):
        return self._cpp_object.masters

    @property
    def coeffs(self):
        return self._cpp_object.coeffs

    @property
    def owners(self):
        return self._cpp_object.owners

    @property
    def offsets(self):
        return self._cpp_object.offsets


class MultiPointConstraint:
    """
    Hold data for multi point constraint relation ships,
    including new index maps for local assembly of matrices and vectors.

    Args:
        V: The function space
        dtype: The dtype of the underlying functions
    """

    _slaves: npt.NDArray[numpy.int32]
    _masters: npt.NDArray[numpy.int64]
    _coeffs: _float_array_types
    _owners: npt.NDArray[numpy.int32]
    _offsets: npt.NDArray[numpy.int32]
    V: _fem.FunctionSpace
    finalized: bool
    _cpp_object: _mpc_classes
    _dtype: numpy.floating
    __slots__ = tuple(__annotations__)

    def __init__(self, V: _fem.FunctionSpace, dtype: numpy.floating = default_scalar_type):
        self._slaves = numpy.array([], dtype=numpy.int32)
        self._masters = numpy.array([], dtype=numpy.int64)
        self._coeffs = numpy.array([], dtype=dtype)  # type: ignore
        self._owners = numpy.array([], dtype=numpy.int32)
        self._offsets = numpy.array([0], dtype=numpy.int32)
        self.V = V
        self.finalized = False
        self._dtype = dtype

    def add_constraint(
        self,
        V: _fem.FunctionSpace,
        slaves: npt.NDArray[numpy.int32],
        masters: npt.NDArray[numpy.int64],
        coeffs: _float_array_types,
        owners: npt.NDArray[numpy.int32],
        offsets: npt.NDArray[numpy.int32],
    ):
        """
        Add new constraint given by numpy arrays.

        Args:
            V: The function space for the constraint
            slaves: List of all slave dofs (using local dof numbering) on this process
            masters: List of all master dofs (using global dof numbering) on this process
            coeffs: The coefficients corresponding to each master.
            owners: The process each master is owned by.
            offsets: Array indicating the location in the masters array for the i-th slave
                in the slaves arrays, i.e.

                .. highlight:: python
                .. code-block:: python

                    masters_of_owned_slave[i] = masters[offsets[i]:offsets[i+1]]

        """
        assert V == self.V
        self._already_finalized()

        if len(slaves) > 0:
            self._offsets = numpy.append(self._offsets, offsets[1:] + len(self._masters))
            self._slaves = numpy.append(self._slaves, slaves)
            self._masters = numpy.append(self._masters, masters)
            self._coeffs = numpy.append(self._coeffs, coeffs)
            self._owners = numpy.append(self._owners, owners)

    def add_constraint_from_mpc_data(self, V: _fem.FunctionSpace, mpc_data: Union[_mpc_data_classes, MPCData]):
        """
        Add new constraint given by an `dolfinc_mpc.cpp.mpc.mpc_data`-object
        """
        self._already_finalized()
        self.add_constraint(
            V,
            mpc_data.slaves,
            mpc_data.masters,
            mpc_data.coeffs,
            mpc_data.owners,
            mpc_data.offsets,
        )

    def finalize(self) -> None:
        """
        Finializes the multi point constraint. After this function is called, no new constraints can be added
        to the constraint. This function creates a map from the cells (local to index) to the slave degrees of
        freedom and builds a new index map and function space where unghosted master dofs are added as ghosts.
        """
        self._already_finalized()
        self._coeffs.astype(numpy.dtype(self._dtype))
        # Initialize C++ object and create slave->cell maps
        if self._dtype == numpy.float32:
            self._cpp_object = dolfinx_mpc.cpp.mpc.MultiPointConstraint_float(
                self.V._cpp_object,
                self._slaves,
                self._masters,
                self._coeffs.astype(self._dtype),
                self._owners,
                self._offsets,
            )
        elif self._dtype == numpy.float64:
            self._cpp_object = dolfinx_mpc.cpp.mpc.MultiPointConstraint_double(
                self.V._cpp_object,
                self._slaves,
                self._masters,
                self._coeffs.astype(self._dtype),
                self._owners,
                self._offsets,
            )
        elif self._dtype == numpy.complex64:
            self._cpp_object = dolfinx_mpc.cpp.mpc.MultiPointConstraint_complex_float(
                self.V._cpp_object,
                self._slaves,
                self._masters,
                self._coeffs.astype(self._dtype),
                self._owners,
                self._offsets,
            )

        elif self._dtype == numpy.complex128:
            self._cpp_object = dolfinx_mpc.cpp.mpc.MultiPointConstraint_complex_double(
                self.V._cpp_object,
                self._slaves,
                self._masters,
                self._coeffs.astype(self._dtype),
                self._owners,
                self._offsets,
            )
        else:
            raise ValueError("Unsupported dtype {coeffs.dtype.type} for coefficients")

        # Replace function space
        self.V = _fem.FunctionSpace(self.V.mesh, self.V.ufl_element(), self._cpp_object.function_space)

        self.finalized = True
        # Delete variables that are no longer required
        del (self._slaves, self._masters, self._coeffs, self._owners, self._offsets)

    def create_periodic_constraint_topological(
        self,
        V: _fem.FunctionSpace,
        meshtag: _mesh.MeshTags,
        tag: int,
        relation: Callable[[numpy.ndarray], numpy.ndarray],
        bcs: List[_fem.DirichletBC],
        scale: _float_classes = default_scalar_type(1.0),
    ):
        """
        Create periodic condition for all closure dofs of on all entities in `meshtag` with value `tag`.
        :math:`u(x_i) = scale * u(relation(x_i))` for all of :math:`x_i` on marked entities.

        Args:
            V: The function space to assign the condition to. Should either be the space of the MPC or a sub space.
               meshtag: MeshTag for entity to apply the periodic condition on
            tag: Tag indicating which entities should be slaves
            relation: Lambda-function describing the geometrical relation
            bcs: Dirichlet boundary conditions for the problem (Periodic constraints will be ignored for these dofs)
            scale: Float for scaling bc
        """
        bcs_ = [bc._cpp_object for bc in bcs]
        if isinstance(scale, numpy.generic):  # nanobind conversion of numpy dtypes to general Python types
            scale = scale.item()  # type: ignore
        if V is self.V:
            mpc_data = dolfinx_mpc.cpp.mpc.create_periodic_constraint_topological(
                self.V._cpp_object, meshtag._cpp_object, tag, relation, bcs_, scale, False
            )
        elif self.V.contains(V):
            mpc_data = dolfinx_mpc.cpp.mpc.create_periodic_constraint_topological(
                V._cpp_object, meshtag._cpp_object, tag, relation, bcs_, scale, True
            )
        else:
            raise RuntimeError("The input space has to be a sub space (or the full space) of the MPC")
        self.add_constraint_from_mpc_data(self.V, mpc_data=mpc_data)

    def create_periodic_constraint_geometrical(
        self,
        V: _fem.FunctionSpace,
        indicator: Callable[[numpy.ndarray], numpy.ndarray],
        relation: Callable[[numpy.ndarray], numpy.ndarray],
        bcs: List[_fem.DirichletBC],
        scale: _float_classes = default_scalar_type(1.0),
    ):
        """
        Create a periodic condition for all degrees of freedom whose physical location satisfies
        :math:`indicator(x_i)==True`, i.e.
        :math:`u(x_i) = scale * u(relation(x_i))` for all :math:`x_i`

        Args:
            V: The function space to assign the condition to. Should either be the space of the MPC or a sub space.
            indicator: Lambda-function to locate degrees of freedom that should be slaves
            relation: Lambda-function describing the geometrical relation to master dofs
            bcs: Dirichlet boundary conditions for the problem
                 (Periodic constraints will be ignored for these dofs)
            scale: Float for scaling bc
        """
        if isinstance(scale, numpy.generic):  # nanobind conversion of numpy dtypes to general Python types
            scale = scale.item()  # type: ignore
        bcs = [] if bcs is None else [bc._cpp_object for bc in bcs]
        if V is self.V:
            mpc_data = dolfinx_mpc.cpp.mpc.create_periodic_constraint_geometrical(
                self.V._cpp_object, indicator, relation, bcs, scale, False
            )
        elif self.V.contains(V):
            mpc_data = dolfinx_mpc.cpp.mpc.create_periodic_constraint_geometrical(
                V._cpp_object, indicator, relation, bcs, scale, True
            )
        else:
            raise RuntimeError("The input space has to be a sub space (or the full space) of the MPC")
        self.add_constraint_from_mpc_data(self.V, mpc_data=mpc_data)

    def create_slip_constraint(
        self,
        space: _fem.FunctionSpace,
        facet_marker: Tuple[_mesh.MeshTags, int],
        v: _fem.Function,
        bcs: List[_fem.DirichletBC] = [],
    ):
        """
        Create a slip constraint :math:`u \\cdot v=0` over the entities defined in `facet_marker` with the given index.

        Args:
            space: Function space (possible sub space) for the current constraint
            facet_marker: Tuple containomg the mesh tag and marker used to locate degrees of freedom
            v: Function containing the directional vector to dot your slip condition (most commonly a normal vector)
            bcs: List of Dirichlet BCs (slip conditions will be ignored on these dofs)

        Examples:
            Create constaint :math:`u\\cdot n=0` of all indices in `mt` marked with `i`

            .. highlight:: python
            .. code-block:: python

                V = dolfinx.fem.functionspace(mesh, ("CG", 1))
                mpc = MultiPointConstraint(V)
                n = dolfinx.fem.Function(V)
                mpc.create_slip_constaint(V, (mt, i), n)

            Create slip constaint for a mixed function space:

            .. highlight:: python
            .. code-block:: python

                cellname = mesh.ufl_cell().cellname()
                Ve = basix.ufl.element(basix.ElementFamily.P, cellname , 2, shape=(mesh.geometry.dim,))
                Qe = basix.ufl.element(basix.ElementFamily.P, cellname , 1)
                me = basix.ufl.mixed_element([Ve, Qe])
                W = dolfinx.fem.functionspace(mesh, me)
                mpc = MultiPointConstraint(W)
                n_space, _ = W.sub(0).collapse()
                normal = dolfinx.fem.Function(n_space)
                mpc.create_slip_constraint(W.sub(0), (mt, i), normal, bcs=[])

            A slip condition cannot be applied on the same degrees of freedom as a Dirichlet BC, and therefore
            any Dirichlet bc for the space of the multi point constraint should be supplied.

            .. highlight:: python
            .. code-block:: python

                cellname = mesh.ufl_cell().cellname()
                Ve = basix.ufl.element(basix.ElementFamily.P, cellname , 2, shape=(mesh.geometry.dim,))
                Qe = basix.ufl.element(basix.ElementFamily.P, cellname , 1)
                me = basix.ufl.mixed_element([Ve, Qe])
                W = dolfinx.fem.functionspace(mesh, me)
                mpc = MultiPointConstraint(W)
                n_space, _ = W.sub(0).collapse()
                normal = Function(n_space)
                bc = dolfinx.fem.dirichletbc(inlet_velocity, dofs, W.sub(0))
                mpc.create_slip_constraint(W.sub(0), (mt, i), normal, bcs=[bc])
        """
        bcs = [] if bcs is None else [bc._cpp_object for bc in bcs]
        if space is self.V:
            sub_space = False
        elif self.V.contains(space):
            sub_space = True
        else:
            raise ValueError("Input space has to be a sub space of the MPC space")
        mpc_data = dolfinx_mpc.cpp.mpc.create_slip_condition(
            space._cpp_object,
            facet_marker[0]._cpp_object,
            facet_marker[1],
            v._cpp_object,
            bcs,
            sub_space,
        )
        self.add_constraint_from_mpc_data(self.V, mpc_data=mpc_data)

    def create_general_constraint(
        self,
        slave_master_dict: Dict[bytes, Dict[bytes, float]],
        subspace_slave: Optional[int] = None,
        subspace_master: Optional[int] = None,
    ):
        """
        Args:
            V: The function space
            slave_master_dict: Nested dictionary, where the first key is the bit representing the slave dof's
                coordinate in the mesh. The item of this key is a dictionary, where each key of this dictionary
                is the bit representation of the master dof's coordinate, and the item the coefficient for
                the MPC equation.
            subspace_slave: If using mixed or vector space, and only want to use dofs from a sub space
                as slave add index here
            subspace_master: Subspace index for mixed or vector spaces

        Example:
            If the dof `D` located at `[d0, d1]` should be constrained to the dofs
            `E` and `F` at `[e0, e1]` and `[f0, f1]` as :math:`D = \\alpha E + \\beta F`
            the dictionary should be:

            .. highlight:: python
            .. code-block:: python

                    {numpy.array([d0, d1], dtype=mesh.geometry.x.dtype).tobytes():
                        {numpy.array([e0, e1], dtype=mesh.geometry.x.dtype).tobytes(): alpha,
                        numpy.array([f0, f1], dtype=mesh.geometry.x.dtype).tobytes(): beta}}
        """
        slaves, masters, coeffs, owners, offsets = create_dictionary_constraint(
            self.V, slave_master_dict, subspace_slave, subspace_master
        )
        self.add_constraint(self.V, slaves, masters, coeffs, owners, offsets)

    def create_contact_slip_condition(
        self,
        meshtags: _mesh.MeshTags,
        slave_marker: int,
        master_marker: int,
        normal: _fem.Function,
        eps2: float = 1e-20,
    ):
        """
        Create a slip condition between two sets of facets marker with individual markers.
        The interfaces should be within machine precision of eachother, but the vertices does not need to align.
        The condition created is :math:`u_s \\cdot normal_s = u_m \\cdot normal_m` where `s` is the
        restriction to the slave facets, `m` to the master facets.

        Args:
            meshtags: The meshtags of the set of facets to tie together
            slave_marker: The marker of the slave facets
            master_marker: The marker of the master facets
            normal: The function used in the dot-product of the constraint
            eps2: The tolerance for the squared distance between cells to be considered as a collision
        """
        if isinstance(eps2, numpy.generic):  # nanobind conversion of numpy dtypes to general Python types
            eps2 = eps2.item()  # type: ignore
        mpc_data = dolfinx_mpc.cpp.mpc.create_contact_slip_condition(
            self.V._cpp_object,
            meshtags._cpp_object,
            slave_marker,
            master_marker,
            normal._cpp_object,
            eps2,
        )
        self.add_constraint_from_mpc_data(self.V, mpc_data)

    def create_contact_inelastic_condition(
        self,
        meshtags: _cpp.mesh.MeshTags_int32,
        slave_marker: int,
        master_marker: int,
        eps2: float = 1e-20,
        allow_missing_masters: bool = False,
    ):
        """
        Create a contact inelastic condition between two sets of facets marker with individual markers.
        The interfaces should be within machine precision of eachother, but the vertices does not need to align.
        The condition created is :math:`u_s = u_m` where `s` is the restriction to the
        slave facets, `m` to the master facets.

        Args:
            meshtags: The meshtags of the set of facets to tie together
            slave_marker: The marker of the slave facets
            master_marker: The marker of the master facets
            eps2: The tolerance for the squared distance between cells to be considered as a collision
            allow_missing_masters: If true, the function will not throw an error if a degree of freedom
                in the closure of the master entities does not have a corresponding set of slave degree
                of freedom.
        """
        if isinstance(eps2, numpy.generic):  # nanobind conversion of numpy dtypes to general Python types
            eps2 = eps2.item()  # type: ignore
        mpc_data = dolfinx_mpc.cpp.mpc.create_contact_inelastic_condition(
            self.V._cpp_object, meshtags._cpp_object, slave_marker, master_marker, eps2, allow_missing_masters
        )
        self.add_constraint_from_mpc_data(self.V, mpc_data)

    @property
    def is_slave(self) -> numpy.ndarray:
        """
        Returns a vector of integers where the ith entry indicates if a degree of freedom (local to process) is a slave.
        """
        self._not_finalized()
        return self._cpp_object.is_slave

    @property
    def slaves(self):
        """
        Returns the degrees of freedom for all slaves local to process
        """
        self._not_finalized()
        return self._cpp_object.slaves

    @property
    def masters(self) -> _cpp.graph.AdjacencyList_int32:
        """
        Returns an adjacency-list whose ith node corresponds to
        a degree of freedom (local to process), and links the corresponding master dofs (local to process).

        Examples:

            .. highlight:: python
            .. code-block:: python

                masters = mpc.masters
                masters_of_dof_i = masters.links(i)
        """
        self._not_finalized()
        return self._cpp_object.masters

    def coefficients(self) -> _float_array_types:
        """
        Returns a vector containing the coefficients for the constraint, and the corresponding offsets
        for the ith degree of freedom.

        Examples:

            .. highlight:: python
            .. code-block:: python

                coeffs, offsets = mpc.coefficients()
                coeffs_of_slave_i = coeffs[offsets[i]:offsets[i+1]]
        """
        self._not_finalized()
        return self._cpp_object.coefficients()

    @property
    def num_local_slaves(self):
        """
        Return the number of slaves owned by the current process.
        """
        self._not_finalized()
        return self._cpp_object.num_local_slaves

    @property
    def cell_to_slaves(self):
        """
        Returns an `dolfinx.cpp.graph.AdjacencyList_int32` whose ith node corresponds to
        the ith cell (local to process), and links the corresponding slave degrees of
        freedom in the cell (local to process).

        Examples:

            .. highlight:: python
            .. code-block:: python

                cell_to_slaves = mpc.cell_to_slaves()
                slaves_in_cell_i = cell_to_slaves.links(i)
        """
        self._not_finalized()
        return self._cpp_object.cell_to_slaves

    @property
    def function_space(self):
        """
        Return the function space for the multi-point constraint with the updated index map
        """
        self._not_finalized()
        return self.V

    def backsubstitution(self, u: Union[_fem.Function, _PETSc.Vec]) -> None:  # type: ignore
        """
        For a Function, impose the multi-point constraint by backsubstiution.
        This function is used after solving the reduced problem to obtain the values
        at the slave degrees of freedom

        .. note::
            It is the users responsibility to destroy the PETSc vector

        Args:
            u: The input function
        """
        try:
            self._cpp_object.backsubstitution(u.x.array)
            u.x.scatter_forward()
        except AttributeError:
            with u.localForm() as vector_local:
                self._cpp_object.backsubstitution(vector_local.array_w)
            u.ghostUpdate(addv=_PETSc.InsertMode.INSERT, mode=_PETSc.ScatterMode.FORWARD)  # type: ignore

    def homogenize(self, u: _fem.Function) -> None:
        """
        For a vector, homogenize (set to zero) the vector components at the multi-point
        constraint slave DoF indices. This is particularly useful for nonlinear problems.

        Args:
            u: The input vector
        """
        self._cpp_object.homogenize(u.x.array)
        u.x.scatter_forward()

    def _already_finalized(self):
        """
        Check if we have already finalized the multi point constraint
        """
        if self.finalized:
            raise RuntimeError("MultiPointConstraint has already been finalized")

    def _not_finalized(self):
        """
        Check if we have finalized the multi point constraint
        """
        if not self.finalized:
            raise RuntimeError("MultiPointConstraint has not been finalized")