# Source code for dolfinx.la.petsc

    # Copyright (C) 2025 Garth N. Wells, Jack S. Hale
    #
    # This file is part of DOLFINx (https://www.fenicsproject.org)
    #
    # SPDX-License-Identifier:    LGPL-3.0-or-later
    """Functions for working with PETSc linear algebra objects.

    Note:
        Due to subtle issues in the interaction between petsc4py memory
        management and the Python garbage collector, it is recommended that
        the PETSc method ``destroy()`` is called on returned PETSc objects
        once the object is no longer required. Note that ``destroy()`` is
        collective over the object's MPI communicator.
    """

    import functools
    import itertools
    import typing
    from collections.abc import Iterable, Sequence

    from petsc4py import PETSc

    import numpy as np
    import numpy.typing as npt

    import dolfinx
    from dolfinx.common import IndexMap
    from dolfinx.la import Vector

    assert dolfinx.has_petsc4py

    __all__ = ["assign", "create_vector", "create_vector_wrap"]

    def _ghost_update(x: PETSc.Vec, insert_mode: PETSc.InsertMode, scatter_mode: PETSc.ScatterMode):  # type: ignore[name-defined]
        """Helper function for ghost updating PETSc vectors"""
        if x.getType() == PETSc.Vec.Type.NEST:  # type: ignore[attr-defined]
            for x_sub in x.getNestSubVecs():
                x_sub.ghostUpdate(addv=insert_mode, mode=scatter_mode)
                x_sub.destroy()
        else:
            x.ghostUpdate(addv=insert_mode, mode=scatter_mode)

    def _zero_vector(x: PETSc.Vec):  # type: ignore[name-defined]
        """Helper function for zeroing out PETSc vectors"""
        if x.getType() == PETSc.Vec.Type.NEST:  # type: ignore[attr-defined]
            for x_sub in x.getNestSubVecs():
                with x_sub.localForm() as x_sub_local:
                    x_sub_local.set(0.0)
                x_sub.destroy()
        else:
            with x.localForm() as x_local:
                x_local.set(0.0)

    def create_vector_wrap(x: Vector) -> PETSc.Vec:  # type: ignore[name-defined]
        """Wrap a distributed DOLFINx vector as a PETSc vector.

        Args:
            x: The vector to wrap as a PETSc vector.

        Returns:
            A PETSc vector that shares data with ``x``.
        """
        index_map = x.index_map
        ghosts = index_map.ghosts.astype(PETSc.IntType)  # type: ignore[attr-defined]
        bs = x.block_size
        size = (index_map.size_local * bs, index_map.size_global * bs)
        return PETSc.Vec().createGhostWithArray(  # type: ignore[attr-defined]
            ghosts, x.array, size=size, bsize=bs, comm=index_map.comm
        )

    def create_vector(
        maps: typing.Sequence[tuple[IndexMap, int]], kind: str | None = None
    ) -> PETSc.Vec:  # type: ignore[name-defined]
        """Create a PETSc vector from a sequence of maps and blocksizes.

        Three cases are supported:

        1. If ``maps=[(im_0, bs_0), ..., (im_n, bs_n)]`` is a sequence of
           indexmaps and blocksizes and ``kind`` is ``None``or is
           ``PETSc.Vec.Type.MPI``, a ghosted PETSc vector whith block structure
           described by ``(im_i, bs_i)`` is created.
           The created vector ``b`` is initialized such that on each MPI
           process ``b = [b_0, b_1, ..., b_n, b_0g, b_1g, ..., b_ng]``, where
           ``b_i`` are the entries associated with the 'owned' degrees-of-
           freedom for ``(im_i, bs_i)`` and ``b_ig`` are the 'unowned' (ghost)
           entries.

           If more than one tuple is supplied, the returned vector has an
           attribute ``_blocks`` that holds the local offsets into ``b`` for
           the (i) owned and (ii) ghost entries for each ``V[i]``. It can be
           accessed by ``b.getAttr("_blocks")``. The offsets can be used to get
           views into ``b`` for blocks, e.g.::

               >>> offsets0, offsets1, = b.getAttr("_blocks")
               >>> offsets0
               (0, 12, 28)
               >>> offsets1
               (28, 32, 35)
               >>> b0_owned = b.array[offsets0[0]:offsets0[1]]
               >>> b0_ghost = b.array[offsets1[0]:offsets1[1]]
               >>> b1_owned = b.array[offsets0[1]:offsets0[2]]
               >>> b1_ghost = b.array[offsets1[1]:offsets1[2]]

        2. If ``V=[(im_0, bs_0), ..., (im_n, bs_n)]`` is a sequence of function
           space and ``kind`` is ``PETSc.Vec.Type.NEST``, a PETSc nested vector
           (a 'nest' of ghosted PETSc vectors) is created.

        Args:
            map: Sequence of tuples of ``IndexMap`` and the associated block
                 size.
            kind: PETSc vector type (``VecType``) to create.

        Returns:
            A PETSc vector with the prescribed layout. The vector is not
            initialised to zero.
        """
        if len(maps) == 1:
            # Single space case
            index_map, bs = maps[0]
            ghosts = index_map.ghosts.astype(PETSc.IntType)  # type: ignore[attr-defined]
            size = (index_map.size_local * bs, index_map.size_global * bs)
            b = PETSc.Vec().createGhost(ghosts, size=size, bsize=bs, comm=index_map.comm)  # type: ignore
            if kind == PETSc.Vec.Type.MPI:  # type: ignore[attr-defined]
                _assign_block_data(maps, b)
            return b

        if kind is None or kind == PETSc.Vec.Type.MPI:  # type: ignore[attr-defined]
            b = dolfinx.cpp.fem.petsc.create_vector_block(maps)
            _assign_block_data(maps, b)
            return b
        elif kind == PETSc.Vec.Type.NEST:  # type: ignore[attr-defined]
            return dolfinx.cpp.fem.petsc.create_vector_nest(maps)
        else:
            raise NotImplementedError(
                "Vector type must be specified for blocked/nested assembly."
                f"Vector type '{kind}' not supported."
                "Did you mean 'nest' or 'mpi'?"
            )

    @functools.singledispatch
    def assign(
        x0: npt.NDArray[np.inexact] | Sequence[npt.NDArray[np.inexact]],
        x1: PETSc.Vec,  # type: ignore[name-defined]
    ):
        """Assign ``x0`` values to a PETSc vector ``x1``.

        Values in ``x0``, which is possibly a stacked collection of arrays,
        are assigned ``x1``. When ``x0`` holds a sequence of ``n``` arrays
        and ``x1`` has type ``NEST``, the assignment is::

                  [x0[0]]
            x1 =  [x0[1]]
                  [.....]
                  [x0[n-1]]

        When ``x0`` holds a sequence of ``n`` arrays and ``x1`` **does
        not** have type ``NEST``, the assignment is::

                  [x0_owned[0]]
            x1 =  [.....]
                  [x0_owned[n-1]]
                  [x0_ghost[0]]
                  [.....]
                  [x0_ghost[n-1]]

        Args:
            x0: An array or list of arrays that will be assigned to ``x1``.
            x1: Vector to assign values to.
        """
        if x1.getType() == PETSc.Vec.Type().NEST:  # type: ignore[attr-defined]
            x1_nest = x1.getNestSubVecs()
            for _x0, _x1 in zip(x0, x1_nest):
                with _x1.localForm() as x:
                    x.array_w[:] = _x0
        else:
            with x1.localForm() as _x:
                if isinstance(x0, Sequence):
                    start = 0
                    for _x0 in x0:
                        end = start + _x0.shape[0]
                        _x.array_w[start:end] = _x0
                        start = end
                else:
                    _x.array_w[:] = x0

    @assign.register
    def _(
        x0: PETSc.Vec,  # type: ignore[name-defined]
        x1: npt.NDArray[np.inexact] | Sequence[npt.NDArray[np.inexact]],
    ):
        """Assign PETSc vector ``x0`` values to (blocked) array(s) ``x1``.

        This function performs the reverse of the assignment performed by
        the version of :func:`.assign(x0: (npt.NDArray[np.inexact] |
        list[npt.NDArray[np.inexact]]), x1: PETSc.Vec)`.

        Args:
            x0: Vector that will have its values assigned to ``x1``.
            x1: An array or list of arrays to assign to.
        """
        if x0.getType() == PETSc.Vec.Type().NEST:  # type: ignore[attr-defined]
            x0_nest = x0.getNestSubVecs()
            for _x0, _x1 in zip(x0_nest, x1):
                with _x0.localForm() as x:
                    _x1[:] = x.array_r[:]  # type: ignore[index]
        else:
            with x0.localForm() as _x0:
                if isinstance(x1, Sequence):
                    start = 0
                    for _x1 in x1:
                        end = start + _x1.shape[0]
                        _x1[:] = _x0.array_r[start:end]
                        start = end
                else:
                    x1[:] = _x0.array_r[:]

    def _assign_block_data(maps: Iterable[tuple[IndexMap, int]], vec: PETSc.Vec):  # type: ignore[name-defined]
        """Assign block data to a PETSc vector.

        Args:
            maps: Iterable of tuples each containing an ``IndexMap`` and the
            associated block size ``bs``.
            vec: PETSc vector to assign block data to.
        """
        # Early exit if the vector already has block data or is a nest vector
        if vec.getAttr("_blocks") is not None or vec.getType() == "nest":
            return

        maps = [(index_map, bs) for index_map, bs in maps]
        off_owned = tuple(
            itertools.accumulate(maps, lambda off, m: off + m[0].size_local * m[1], initial=0)
        )
        off_ghost = tuple(
            itertools.accumulate(
                maps, lambda off, m: off + m[0].num_ghosts * m[1], initial=off_owned[-1]
            )
        )
        vec.setAttr("_blocks", (off_owned, off_ghost))