from __future__ import annotations

from typing import TYPE_CHECKING, Protocol, Self

if TYPE_CHECKING:
    import h5py
    import numpy as np
    from numpy.typing import NDArray


class DataIndex(Protocol):
    """
    A DataIndex is responsible for tracking which elements of some hdf5 dataset
    should be read. Note these are NOT a spatial index, though a spatial index
    does return a DataIndex when it is queried.
    """

    @classmethod
    def from_size(cls, size: int) -> Self: ...
    def get_data(self, data: h5py.Dataset | np.ndarray) -> np.ndarray: ...
    def n_in_range(
        self, start: NDArray[np.int_], end: NDArray[np.int_]
    ) -> NDArray[np.int_]: ...
    def intersection(self, other: DataIndex) -> DataIndex: ...
    def into_mask(self) -> NDArray[np.bool_]: ...
    def into_array(self) -> NDArray[np.int_]: ...
    def mask(self, mask: np.ndarray) -> DataIndex: ...
    def range(self) -> tuple[int, int]: ...
    def concatenate(self, *others: DataIndex) -> DataIndex: ...
    def __len__(self) -> int: ...
    def __getitem__(self, item: int) -> DataIndex: ...
