from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any

import numpy as np
from numpy.lib.mixins import NDArrayOperatorsMixin
from numpy.typing import ArrayLike, DTypeLike, NDArray
from typing_extensions import Self, override

@dataclass(frozen=True, slots=True)
class Feature(NDArrayOperatorsMixin, ABC):
    value: NDArray[Any]
    def __init__(self, value: NDArray[Any]) -> None: ...
    __array_priority__: int
    @property
    def shape(self) -> tuple[int, ...]: ...
    @property
    def ndim(self) -> int: ...
    @property
    def dtype(self) -> DTypeLike: ...
    @property
    def size(self) -> int: ...
    def mean(self, axis: int | None = None, **kwargs: Any) -> Any: ...  # noqa: ANN401
    def sum(self, axis: int | None = None, **kwargs: Any) -> Any: ...  # noqa: ANN401
    def min(self, axis: int | None = None, **kwargs: Any) -> Any: ...  # noqa: ANN401
    def max(self, axis: int | None = None, **kwargs: Any) -> Any: ...  # noqa: ANN401
    @abstractmethod
    def crop(self, indices: NDArray[np.integer]) -> Self: ...
    @abstractmethod
    def copy(self) -> Self: ...
    @abstractmethod
    def __getitem__(self, key: Any) -> Self: ...  # noqa: ANN401
    def __len__(self) -> int: ...
    def __bool__(self) -> bool: ...
    def __array__(self, dtype: DTypeLike | None = None) -> NDArray[Any]: ...
    def __array_ufunc__(
        self,
        ufunc: np.ufunc,
        method: str,
        *inputs: Any,  # noqa: ANN401
        **kwargs: Any,  # noqa: ANN401
    ) -> Any: ...  # noqa: ANN401
    def __add__(self, other: ArrayLike) -> Self: ...
    def __radd__(self, other: ArrayLike) -> Self: ...
    def __sub__(self, other: ArrayLike) -> Self: ...
    def __rsub__(self, other: ArrayLike) -> Self: ...
    def __mul__(self, other: ArrayLike) -> Self: ...
    def __rmul__(self, other: ArrayLike) -> Self: ...
    def __truediv__(self, other: ArrayLike) -> Self: ...
    def __rtruediv__(self, other: ArrayLike) -> Self: ...
    def __floordiv__(self, other: ArrayLike) -> Self: ...
    def __rfloordiv__(self, other: ArrayLike) -> Self: ...
    def __mod__(self, other: ArrayLike) -> Self: ...
    def __rmod__(self, other: ArrayLike) -> Self: ...
    def __pow__(self, other: ArrayLike) -> Self: ...
    def __rpow__(self, other: ArrayLike) -> Self: ...
    def __matmul__(self, other: ArrayLike) -> Self: ...
    def __rmatmul__(self, other: ArrayLike) -> Self: ...
    def __neg__(self) -> Self: ...
    def __pos__(self) -> Self: ...
    def __abs__(self) -> Self: ...
    def __invert__(self) -> Self: ...
    def __and__(self, other: ArrayLike) -> NDArray[np.bool_]: ...
    def __rand__(self, other: ArrayLike) -> NDArray[np.bool_]: ...
    def __or__(self, other: ArrayLike) -> NDArray[np.bool_]: ...
    def __ror__(self, other: ArrayLike) -> NDArray[np.bool_]: ...
    def __xor__(self, other: ArrayLike) -> NDArray[np.bool_]: ...
    def __rxor__(self, other: ArrayLike) -> NDArray[np.bool_]: ...
    def __lshift__(self, other: ArrayLike) -> NDArray[np.integer]: ...
    def __rlshift__(self, other: ArrayLike) -> NDArray[np.integer]: ...
    def __rshift__(self, other: ArrayLike) -> NDArray[np.integer]: ...
    def __rrshift__(self, other: ArrayLike) -> NDArray[np.integer]: ...
    def __lt__(self, other: ArrayLike) -> NDArray[np.bool_]: ...
    def __le__(self, other: ArrayLike) -> NDArray[np.bool_]: ...
    def __gt__(self, other: ArrayLike) -> NDArray[np.bool_]: ...
    def __ge__(self, other: ArrayLike) -> NDArray[np.bool_]: ...
    def __eq__(self, other: ArrayLike) -> NDArray[np.bool_]: ...  # pyright: ignore[reportIncompatibleMethodOverride]

class NodeFeature(Feature):
    @override
    def crop(self, indices: NDArray[np.integer]) -> Self: ...
    @override
    def copy(self) -> Self: ...
    @override
    def __getitem__(self, key: Any) -> Self: ...

class EdgeFeature(Feature):
    src_indices: NDArray[np.integer]
    dst_indices: NDArray[np.integer]
    def __init__(
        self,
        value: NDArray[Any],
        src_indices: NDArray[np.integer] = ...,
        dst_indices: NDArray[np.integer] = ...,
    ) -> None: ...
    def __post_init__(self) -> None: ...
    @property
    def src(self) -> NDArray[np.integer]: ...
    @property
    def dst(self) -> NDArray[np.integer]: ...
    @property
    def nodes(self) -> NDArray[np.integer]: ...
    @override
    def crop(self, indices: NDArray[np.integer]) -> Self: ...
    @override
    def copy(self) -> Self: ...
    @override
    def __getitem__(self, key: Any) -> Self: ...
