from collections.abc import Iterator
from typing import (
    Any,
    Literal,
    Protocol,
    SupportsIndex,
    overload,
    runtime_checkable,
)
from typing_extensions import Self

import numpy as np
import numpy.typing as npt

from ._bsr import bsr_matrix
from ._coo import coo_matrix
from ._csc import csc_matrix
from ._csr import csr_matrix
from ._dia import dia_matrix
from ._dok import dok_matrix
from ._lil import lil_matrix
from .typing import (
    _SCT,
    SparseArray,
    _ArrayType,
    _CastingKind,
    _DType_co,
    _DTypeLike,
    _Formats,
    _IntLike_co,
    _NumberLike_co,
    _OrderType,
    _SCT_uifcO,
    _ShapeAnno,
    _ShapeLike,
)

# Although spmatrix is not a Protocol, I hint it as such such that if
# subclassed it is clear that the methods and attributes are meant to be
# implemented. All methods and attributes are as in the private base class
# _spbase (except for the missing `def __init__`)
@runtime_checkable
class spmatrix(Protocol[_ShapeAnno, _DType_co]):
{cls_body}
    # spmatrix format can be any value of its subclasses, which I leave as str
    @property
    def format(self) -> str: ...

    ###########################################################################
    # _matrix methods that override _spbase def
    ###########################################################################
    def __mul__(
        self, other: npt.NDArray[Any] | SparseArray[Any]
    ) -> npt.NDArray[Any] | SparseArray[Any]: ...
    def __rmul__(
        self, other: npt.NDArray[Any] | SparseArray[Any]
    ) -> npt.NDArray[Any] | SparseArray[Any]: ...
    def __pow__(self, power: _IntLike_co) -> spmatrix[Any, Any]: ...

    ###########################################################################
    # spmatrix only methods
    ###########################################################################
    def set_shape(self, shape: tuple[SupportsIndex, SupportsIndex]) -> None: ...
    def get_shape(self) -> tuple[int, int]: ...
    # The shape.setter is not annotated altough it is defined
    def asfptype(self) -> spmatrix[Any, Any]: ...
    def getmaxprint(self) -> int: ...
    def getformat(self) -> _Formats: ...
    @overload
    def getnnz(self, axis: None = ...) -> int: ...
    @overload
    def getnnz(self, axis: SupportsIndex) -> npt.NDArray[np.int_]: ...
    def getH(self) -> spmatrix[Any, _DType_co]: ...
    def getcol(self, j: SupportsIndex) -> spmatrix[Any, _DType_co]: ...
    def getrow(self, i: SupportsIndex) -> spmatrix[Any, _DType_co]: ...
