from __future__ import annotations

import numpy as np
import pyarrow as pa

from lenskit.data.matrix import SparseRowArray
from lenskit.data.types import ID

def transpose_csr(
    matrix: SparseRowArray, permute: bool
) -> tuple[pa.Array, pa.Int32Array, pa.Array | None]: ...
def is_sorted_coo(data: list[pa.RecordBatch], c1: str, c2: str) -> bool: ...
def argsort_descending(data: pa.Array) -> pa.Int32Array: ...
def negative_mask(n: int, indices: pa.Int32Array) -> pa.BooleanArray:
    """
    Efficiently create a boolean array that is true everywhere except ``indices``.
    """
    ...

def sample_negatives(
    coords: CoordinateTable,
    rows: np.ndarray[tuple[int], np.dtype[np.int32]],
    n_cols: int,
    *,
    n: int = 1,
    max_attempts: int = 10,
    pop_weighted: bool = False,
    seed: int,
) -> np.ndarray[tuple[int, int], np.dtype[np.int32]]: ...
def hash_array(arr: pa.Array, /) -> str: ...

class IDIndex:
    """
    Backend implementation for ID indexes / vocabularies.
    """
    def __init__(self, ids: pa.Array | None = None): ...
    def get_index(self, id: ID): ...
    def get_indexes(self, ids: pa.Array): ...
    def id_array(self) -> pa.Array: ...
    def __len__(self) -> int: ...

class CoordinateTable:
    def __init__(self, dims: int): ...
    def copy(self) -> CoordinateTable: ...
    def dimensions(self) -> int: ...
    def __len__(self) -> int: ...
    def entry_count(self) -> int: ...
    def unique_count(self) -> int: ...
    def extend(self, batch: pa.RecordBatch | list[pa.RecordBatch]) -> tuple[int, int]: ...
    def contains(self, *coords: int) -> bool: ...
    def find(self, *coords: int) -> int | None: ...
    def _get_coord(self, dim: int, position: int) -> int: ...
