from typing import Protocol, Self

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 take(self, n: int, at: str = "random") -> "DataIndex": ...
    def take_range(self, start: int, end: int) -> "DataIndex": ...
    def intersection(self, other: "DataIndex") -> "DataIndex": ...
    def projection(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": ...
