# Copyright (C) 2017-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
"""Linear algebra functionality."""

from typing import Generic, TypeVar

import numpy as np
import numpy.typing as npt

import dolfinx
from dolfinx import cpp as _cpp
from dolfinx.cpp.common import IndexMap
from dolfinx.cpp.la import BlockMode, InsertMode, Norm
from dolfinx.typing import Scalar

__all__ = [
    "InsertMode",
    "MatrixCSR",
    "Norm",
    "Vector",
    "is_orthonormal",
    "matrix_csr",
    "norm",
    "orthonormalize",
    "vector",
]


_T = TypeVar("_T", np.float32, np.float64, np.complex64, np.complex128, np.int8, np.int32, np.int64)


class Vector(Generic[_T]):
    """Distributed vector object."""

    _cpp_object: (
        _cpp.la.Vector_float32
        | _cpp.la.Vector_float64
        | _cpp.la.Vector_complex64
        | _cpp.la.Vector_complex128
        | _cpp.la.Vector_int8
        | _cpp.la.Vector_int32
        | _cpp.la.Vector_int64
    )

    def __init__(
        self,
        x: (
            _cpp.la.Vector_float32
            | _cpp.la.Vector_float64
            | _cpp.la.Vector_complex64
            | _cpp.la.Vector_complex128
            | _cpp.la.Vector_int8
            | _cpp.la.Vector_int32
            | _cpp.la.Vector_int64
        ),
    ):
        """Create a distributed vector.

        Args:
            x: C++ Vector object.

        Note:
            This initialiser is intended for internal library use only.
            User code should call :func:`vector` to create a vector object.
        """
        self._cpp_object = x
        self._petsc_x = None

    def __del__(self):
        """Delete the PETSc vector if it was created."""
        if self._petsc_x is not None:
            self._petsc_x.destroy()

    @property
    def index_map(self) -> IndexMap:
        """Index map that describes size and parallel distribution."""
        return self._cpp_object.index_map

    @property
    def block_size(self) -> int:
        """Block size for the vector."""
        return self._cpp_object.bs

    @property
    def array(self) -> npt.NDArray[_T]:
        """Local representation of the vector."""
        return self._cpp_object.array

    @property
    def petsc_vec(self):
        """PETSc vector holding the entries of the vector.

        Upon first call, this function creates a PETSc ``Vec`` object
        that wraps the degree-of-freedom data. The ``Vec`` object is
        cached and the cached ``Vec`` is returned upon subsequent calls.

        Note:
          When the object is destroyed it will destroy the underlying
          petsc4py vector automatically.
        """
        assert dolfinx.has_petsc4py

        from dolfinx.la.petsc import create_vector_wrap

        if self._petsc_x is None:
            self._petsc_x = create_vector_wrap(self)
        return self._petsc_x

    def scatter_forward(self) -> None:
        """Update ghost entries."""
        self._cpp_object.scatter_forward()

    def scatter_reverse(self, mode: InsertMode) -> None:
        """Scatter ghost entries to owner.

        Args:
            mode: Control how scattered values are set/accumulated by
                owner.
        """
        self._cpp_object.scatter_reverse(mode)


class MatrixCSR(Generic[Scalar]):
    """Distributed compressed sparse row matrix."""

    _cpp_object: (
        _cpp.la.MatrixCSR_float32
        | _cpp.la.MatrixCSR_float64
        | _cpp.la.MatrixCSR_complex64
        | _cpp.la.MatrixCSR_complex128
    )

    def __init__(
        self,
        A: (
            _cpp.la.MatrixCSR_float32
            | _cpp.la.MatrixCSR_float64
            | _cpp.la.MatrixCSR_complex64
            | _cpp.la.MatrixCSR_complex128
        ),
    ):
        """Create a distributed compressed sparse row matrix.

        Note:
            Objects of this type should be created using
            :func:`matrix_csr` and not created using this initialiser.

        Args:
            A: The C++/nanobind matrix object.
        """
        self._cpp_object = A

    def index_map(self, i: int) -> IndexMap:
        """Index map for row/column.

        Args:
            i: 0 for row map, 1 for column map.
        """
        return self._cpp_object.index_map(i)

    def mult(self, x: Vector[Scalar], y: Vector[Scalar], transpose: bool = False) -> None:
        """Compute ``y += Ax`` or ``y += A^T x``.

        Args:
            x: Input Vector
            y: Output Vector
            transpose: if True, compute y += A^T x
        """
        if transpose:
            self._cpp_object.multT(x._cpp_object, y._cpp_object)
        else:
            self._cpp_object.mult(x._cpp_object, y._cpp_object)

    def matmul(self, B):
        """Compute matrix product ``A * B``, where `A` is this matrix.

        Args:
            B: Input Matrix to multiply by
        """
        if (
            self.index_map(1).size_local != B.index_map(0).size_local
            or self.index_map(1).size_global != B.index_map(0).size_global
        ):
            raise RuntimeError("Invalid matrix sizes for matmul.")
        if (
            self.block_size[0] != 1
            or self.block_size[1] != 1
            or B.block_size[0] != 1
            or B.block_size[1] != 1
        ):
            raise RuntimeError("Block size not supported in matmul.")

        return MatrixCSR(self._cpp_object.mult(B._cpp_object))

    def transpose(self):
        """Compute transpose matrix."""
        return MatrixCSR(self._cpp_object.transpose())

    @property
    def block_size(self) -> list:
        """Block sizes for the matrix."""
        return self._cpp_object.bs

    def add(
        self,
        x: npt.NDArray[Scalar],
        rows: npt.NDArray[np.int32],
        cols: npt.NDArray[np.int32],
        bs: int = 1,
    ) -> None:
        """Add a block of values in the matrix."""
        self._cpp_object.add(x, rows, cols, bs)

    def set(
        self,
        x: npt.NDArray[Scalar],
        rows: npt.NDArray[np.int32],
        cols: npt.NDArray[np.int32],
        bs: int = 1,
    ) -> None:
        """Set a block of values in the matrix."""
        self._cpp_object.set(x, rows, cols, bs)

    def set_value(self, x: Scalar) -> None:
        """Set all non-zero entries to a value.

        Args:
            x: The value to set all non-zero entries to.
        """
        self.data[:] = x

    def scatter_reverse(self) -> None:
        """Scatter and accumulate ghost values."""
        self._cpp_object.scatter_reverse()

    def squared_norm(self) -> np.floating:
        """Compute the squared Frobenius norm.

        Note:
            This operation is collective and requires communication.
        """
        return self._cpp_object.squared_norm()

    @property
    def data(self) -> npt.NDArray[Scalar]:
        """Underlying matrix entry data."""
        return self._cpp_object.data

    @property
    def indices(self) -> npt.NDArray[np.int32]:
        """Local column indices."""
        return self._cpp_object.indices

    @property
    def indptr(self) -> npt.NDArray[np.int64]:
        """Local row pointers."""
        return self._cpp_object.indptr

    def to_dense(self) -> npt.NDArray[Scalar]:
        """Copy to a dense 2D array.

        Note:
            Typically used for debugging.
        """
        return self._cpp_object.to_dense()

    def to_scipy(self, ghosted: bool = False):
        """Convert to a SciPy CSR/BSR matrix. Data is shared.

        Note:
            SciPy must be available.

        Args:
            ghosted: If ``True`` rows that are ghosted in parallel are
                included in the returned SciPy matrix, otherwise ghost
                rows are not included.

        Returns:
            SciPy compressed sparse row (both block sizes equal to one)
            or a SciPy block compressed sparse row matrix.
        """
        bs0, bs1 = self._cpp_object.bs
        ncols = self.index_map(1).size_local + self.index_map(1).num_ghosts
        if ghosted:
            nrows = self.index_map(0).size_local + self.index_map(0).num_ghosts
            data, indices, indptr = self.data, self.indices, self.indptr
        else:
            nrows = self.index_map(0).size_local
            nnzlocal = self.indptr[nrows]
            data, indices, indptr = (
                self.data[: (bs0 * bs1) * nnzlocal],
                self.indices[:nnzlocal],
                self.indptr[: nrows + 1],
            )

        if bs0 == 1 and bs1 == 1:
            from scipy.sparse import csr_matrix as _csr

            return _csr((data, indices, indptr), shape=(nrows, ncols))
        else:
            from scipy.sparse import bsr_matrix as _bsr

            return _bsr(
                (data.reshape(-1, bs0, bs1), indices, indptr), shape=(bs0 * nrows, bs1 * ncols)
            )


def matrix_csr(
    sp: _cpp.la.SparsityPattern, block_mode=BlockMode.compact, dtype: npt.DTypeLike = np.float64
) -> MatrixCSR:
    """Create a distributed sparse matrix.

    The matrix uses compressed sparse row storage.

    Args:
        sp: The sparsity pattern that defines the nonzero structure of
            the matrix the parallel distribution of the matrix.
        block_mode: Block mode to use.
        dtype: Scalar type.

    Returns:
        A sparse matrix.
    """
    if np.issubdtype(dtype, np.float32):
        ftype = _cpp.la.MatrixCSR_float32
    elif np.issubdtype(dtype, np.float64):
        ftype = _cpp.la.MatrixCSR_float64
    elif np.issubdtype(dtype, np.complex64):
        ftype = _cpp.la.MatrixCSR_complex64
    elif np.issubdtype(dtype, np.complex128):
        ftype = _cpp.la.MatrixCSR_complex128
    else:
        raise NotImplementedError(f"Type {dtype} not supported.")

    return MatrixCSR(ftype(sp, block_mode))


def vector(map, bs=1, dtype: npt.DTypeLike = np.float64) -> Vector:
    """Create a distributed vector.

    Args:
        map: Index map the describes the size and distribution of the
            vector.
        bs: Block size.
        dtype: The scalar type.

    Returns:
        A distributed vector.
    """
    if np.issubdtype(dtype, np.float32):
        vtype = _cpp.la.Vector_float32
    elif np.issubdtype(dtype, np.float64):
        vtype = _cpp.la.Vector_float64
    elif np.issubdtype(dtype, np.complex64):
        vtype = _cpp.la.Vector_complex64
    elif np.issubdtype(dtype, np.complex128):
        vtype = _cpp.la.Vector_complex128
    elif np.issubdtype(dtype, np.int8):
        vtype = _cpp.la.Vector_int8
    elif np.issubdtype(dtype, np.int32):
        vtype = _cpp.la.Vector_int32
    elif np.issubdtype(dtype, np.int64):
        vtype = _cpp.la.Vector_int64
    else:
        raise NotImplementedError(f"Type {dtype} not supported.")

    return Vector(vtype(map, bs))


def orthonormalize(basis: list[Vector[_T]]) -> None:
    """Orthogonalise set of vectors in-place."""
    _cpp.la.orthonormalize([x._cpp_object for x in basis])


def is_orthonormal(basis: list[Vector[_T]], eps: float = 1.0e-12) -> bool:
    """Check that list of vectors are orthonormal."""
    return _cpp.la.is_orthonormal([x._cpp_object for x in basis], eps)


def norm(x: Vector[_T], type: _cpp.la.Norm = _cpp.la.Norm.l2) -> np.floating:
    """Compute a norm of the vector.

    Args:
        x: Vector to measure.
        type: Norm type to compute.

    Returns:
        Computed norm.
    """
    return _cpp.la.norm(x._cpp_object, type)
