from _typeshed import Incomplete
from collections.abc import Callable, Iterable, Mapping, Sequence
from typing import Any, Literal, overload

import numpy as np
import numpy.typing as npt
from numpy import _CastingKind
from numpy._typing import _AnyShape, _DTypeLike, _DTypeLikeVoid, _Shape
from numpy.ma.mrecords import MaskedRecords

__all__ = [
    "append_fields",
    "apply_along_fields",
    "assign_fields_by_name",
    "drop_fields",
    "find_duplicates",
    "flatten_descr",
    "get_fieldstructure",
    "get_names",
    "get_names_flat",
    "join_by",
    "merge_arrays",
    "rec_append_fields",
    "rec_drop_fields",
    "rec_join",
    "recursive_fill_fields",
    "rename_fields",
    "repack_fields",
    "require_fields",
    "stack_arrays",
    "structured_to_unstructured",
    "unstructured_to_structured",
]

type _OneOrMany[T] = T | Iterable[T]
type _BuiltinSequence[T] = tuple[T, ...] | list[T]

type _NestedNames = tuple[str | _NestedNames, ...]
type _NonVoid = np.bool | np.number | np.character | np.datetime64 | np.timedelta64 | np.object_
type _NonVoidDType = np.dtype[_NonVoid] | np.dtypes.StringDType

type _JoinType = Literal["inner", "outer", "leftouter"]

###

def recursive_fill_fields[VoidArrayT: npt.NDArray[np.void]](input: npt.NDArray[np.void], output: VoidArrayT) -> VoidArrayT: ...

#
def get_names(adtype: np.dtype[np.void]) -> _NestedNames: ...
def get_names_flat(adtype: np.dtype[np.void]) -> tuple[str, ...]: ...

#
@overload
def flatten_descr[NonVoidDTypeT: _NonVoidDType](ndtype: NonVoidDTypeT) -> tuple[tuple[Literal[""], NonVoidDTypeT]]: ...
@overload
def flatten_descr(ndtype: np.dtype[np.void]) -> tuple[tuple[str, np.dtype]]: ...

#
def get_fieldstructure(
    adtype: np.dtype[np.void],
    lastname: str | None = None,
    parents: dict[str, list[str]] | None = None,
) -> dict[str, list[str]]: ...

#
@overload
def merge_arrays[ShapeT: _Shape](
    seqarrays: Sequence[np.ndarray[ShapeT, np.dtype]] | np.ndarray[ShapeT, np.dtype],
    fill_value: float = -1,
    flatten: bool = False,
    usemask: bool = False,
    asrecarray: bool = False,
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...
@overload
def merge_arrays(
    seqarrays: Sequence[npt.ArrayLike] | np.void,
    fill_value: float = -1,
    flatten: bool = False,
    usemask: bool = False,
    asrecarray: bool = False,
) -> np.recarray[_AnyShape, np.dtype[np.void]]: ...

#
@overload
def drop_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    drop_names: str | Iterable[str],
    usemask: bool = True,
    asrecarray: Literal[False] = False,
) -> np.ndarray[ShapeT, np.dtype[np.void]]: ...
@overload
def drop_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    drop_names: str | Iterable[str],
    usemask: bool,
    asrecarray: Literal[True],
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...
@overload
def drop_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    drop_names: str | Iterable[str],
    usemask: bool = True,
    *,
    asrecarray: Literal[True],
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...

#
@overload
def rename_fields[ShapeT: _Shape](
    base: MaskedRecords[ShapeT, np.dtype[np.void]],
    namemapper: Mapping[str, str],
) -> MaskedRecords[ShapeT, np.dtype[np.void]]: ...
@overload
def rename_fields[ShapeT: _Shape](
    base: np.ma.MaskedArray[ShapeT, np.dtype[np.void]],
    namemapper: Mapping[str, str],
) -> np.ma.MaskedArray[ShapeT, np.dtype[np.void]]: ...
@overload
def rename_fields[ShapeT: _Shape](
    base: np.recarray[ShapeT, np.dtype[np.void]],
    namemapper: Mapping[str, str],
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...
@overload
def rename_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    namemapper: Mapping[str, str],
) -> np.ndarray[ShapeT, np.dtype[np.void]]: ...

#
@overload
def append_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    names: _OneOrMany[str],
    data: _OneOrMany[npt.NDArray[Any]],
    dtypes: _BuiltinSequence[np.dtype] | None,
    fill_value: int,
    usemask: Literal[False],
    asrecarray: Literal[False] = False,
) -> np.ndarray[ShapeT, np.dtype[np.void]]: ...
@overload
def append_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    names: _OneOrMany[str],
    data: _OneOrMany[npt.NDArray[Any]],
    dtypes: _BuiltinSequence[np.dtype] | None = None,
    fill_value: int = -1,
    *,
    usemask: Literal[False],
    asrecarray: Literal[False] = False,
) -> np.ndarray[ShapeT, np.dtype[np.void]]: ...
@overload
def append_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    names: _OneOrMany[str],
    data: _OneOrMany[npt.NDArray[Any]],
    dtypes: _BuiltinSequence[np.dtype] | None,
    fill_value: int,
    usemask: Literal[False],
    asrecarray: Literal[True],
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...
@overload
def append_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    names: _OneOrMany[str],
    data: _OneOrMany[npt.NDArray[Any]],
    dtypes: _BuiltinSequence[np.dtype] | None = None,
    fill_value: int = -1,
    *,
    usemask: Literal[False],
    asrecarray: Literal[True],
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...
@overload
def append_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    names: _OneOrMany[str],
    data: _OneOrMany[npt.NDArray[Any]],
    dtypes: _BuiltinSequence[np.dtype] | None = None,
    fill_value: int = -1,
    usemask: Literal[True] = True,
    asrecarray: Literal[False] = False,
) -> np.ma.MaskedArray[ShapeT, np.dtype[np.void]]: ...
@overload
def append_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    names: _OneOrMany[str],
    data: _OneOrMany[npt.NDArray[Any]],
    dtypes: _BuiltinSequence[np.dtype] | None,
    fill_value: int,
    usemask: Literal[True],
    asrecarray: Literal[True],
) -> MaskedRecords[ShapeT, np.dtype[np.void]]: ...
@overload
def append_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    names: _OneOrMany[str],
    data: _OneOrMany[npt.NDArray[Any]],
    dtypes: _BuiltinSequence[np.dtype] | None = None,
    fill_value: int = -1,
    usemask: Literal[True] = True,
    *,
    asrecarray: Literal[True],
) -> MaskedRecords[ShapeT, np.dtype[np.void]]: ...

#
def rec_drop_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    drop_names: str | Iterable[str],
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...

#
def rec_append_fields[ShapeT: _Shape](
    base: np.ndarray[ShapeT, np.dtype[np.void]],
    names: _OneOrMany[str],
    data: _OneOrMany[npt.NDArray[Any]],
    dtypes: _BuiltinSequence[np.dtype] | None = None,
) -> np.ma.MaskedArray[ShapeT, np.dtype[np.void]]: ...

# TODO(jorenham): Stop passing `void` directly once structured dtypes are implemented,
# e.g. using a `TypeVar` with constraints.
# https://github.com/numpy/numtype/issues/92
@overload
def repack_fields[DTypeT: np.dtype](a: DTypeT, align: bool = False, recurse: bool = False) -> DTypeT: ...
@overload
def repack_fields[ScalarT: np.generic](a: ScalarT, align: bool = False, recurse: bool = False) -> ScalarT: ...
@overload
def repack_fields[ArrayT: np.ndarray](a: ArrayT, align: bool = False, recurse: bool = False) -> ArrayT: ...

# TODO(jorenham): Attempt shape-typing (return type has ndim == arr.ndim + 1)
@overload
def structured_to_unstructured[ScalarT: np.generic](
    arr: npt.NDArray[np.void],
    dtype: _DTypeLike[ScalarT],
    copy: bool = False,
    casting: _CastingKind = "unsafe",
) -> npt.NDArray[ScalarT]: ...
@overload
def structured_to_unstructured(
    arr: npt.NDArray[np.void],
    dtype: npt.DTypeLike | None = None,
    copy: bool = False,
    casting: _CastingKind = "unsafe",
) -> npt.NDArray[Any]: ...

#
@overload
def unstructured_to_structured(
    arr: npt.NDArray[Any],
    dtype: npt.DTypeLike,
    names: None = None,
    align: bool = False,
    copy: bool = False,
    casting: str = "unsafe",
) -> npt.NDArray[np.void]: ...
@overload
def unstructured_to_structured(
    arr: npt.NDArray[Any],
    dtype: None,
    names: _OneOrMany[str],
    align: bool = False,
    copy: bool = False,
    casting: str = "unsafe",
) -> npt.NDArray[np.void]: ...
@overload
def unstructured_to_structured(
    arr: npt.NDArray[Any],
    dtype: None = None,
    *,
    names: _OneOrMany[str],
    align: bool = False,
    copy: bool = False,
    casting: str = "unsafe",
) -> npt.NDArray[np.void]: ...

#
def apply_along_fields[ShapeT: _Shape](
    func: Callable[[np.ndarray[ShapeT]], np.ndarray],
    arr: np.ndarray[ShapeT, np.dtype[np.void]],
) -> np.ndarray[ShapeT, np.dtype[np.void]]: ...

#
def assign_fields_by_name(dst: npt.NDArray[np.void], src: npt.NDArray[np.void], zero_unassigned: bool = True) -> None: ...

#
def require_fields[ShapeT: _Shape](
    array: np.ndarray[ShapeT, np.dtype[np.void]],
    required_dtype: _DTypeLikeVoid,
) -> np.ndarray[ShapeT, np.dtype[np.void]]: ...

# TODO(jorenham): Attempt shape-typing
@overload
def stack_arrays[ArrayT: np.ndarray](
    arrays: ArrayT,
    defaults: Mapping[str, object] | None = None,
    usemask: bool = True,
    asrecarray: bool = False,
    autoconvert: bool = False,
) -> ArrayT: ...
@overload
def stack_arrays(
    arrays: Sequence[npt.NDArray[Any]],
    defaults: Mapping[str, Incomplete] | None,
    usemask: Literal[False],
    asrecarray: Literal[False] = False,
    autoconvert: bool = False,
) -> npt.NDArray[np.void]: ...
@overload
def stack_arrays(
    arrays: Sequence[npt.NDArray[Any]],
    defaults: Mapping[str, Incomplete] | None = None,
    *,
    usemask: Literal[False],
    asrecarray: Literal[False] = False,
    autoconvert: bool = False,
) -> npt.NDArray[np.void]: ...
@overload
def stack_arrays(
    arrays: Sequence[npt.NDArray[Any]],
    defaults: Mapping[str, Incomplete] | None = None,
    *,
    usemask: Literal[False],
    asrecarray: Literal[True],
    autoconvert: bool = False,
) -> np.recarray[_AnyShape, np.dtype[np.void]]: ...
@overload
def stack_arrays(
    arrays: Sequence[npt.NDArray[Any]],
    defaults: Mapping[str, Incomplete] | None = None,
    usemask: Literal[True] = True,
    asrecarray: Literal[False] = False,
    autoconvert: bool = False,
) -> np.ma.MaskedArray[_AnyShape, np.dtype[np.void]]: ...
@overload
def stack_arrays(
    arrays: Sequence[npt.NDArray[Any]],
    defaults: Mapping[str, Incomplete] | None,
    usemask: Literal[True],
    asrecarray: Literal[True],
    autoconvert: bool = False,
) -> MaskedRecords[_AnyShape, np.dtype[np.void]]: ...
@overload
def stack_arrays(
    arrays: Sequence[npt.NDArray[Any]],
    defaults: Mapping[str, Incomplete] | None = None,
    usemask: Literal[True] = True,
    *,
    asrecarray: Literal[True],
    autoconvert: bool = False,
) -> MaskedRecords[_AnyShape, np.dtype[np.void]]: ...

#
@overload
def find_duplicates[ShapeT: _Shape](
    a: np.ma.MaskedArray[ShapeT, np.dtype[np.void]],
    key: str | None = None,
    ignoremask: bool = True,
    return_index: Literal[False] = False,
) -> np.ma.MaskedArray[ShapeT, np.dtype[np.void]]: ...
@overload
def find_duplicates[ShapeT: _Shape](
    a: np.ma.MaskedArray[ShapeT, np.dtype[np.void]],
    key: str | None,
    ignoremask: bool,
    return_index: Literal[True],
) -> tuple[np.ma.MaskedArray[ShapeT, np.dtype[np.void]], np.ndarray[ShapeT, np.dtype[np.int_]]]: ...
@overload
def find_duplicates[ShapeT: _Shape](
    a: np.ma.MaskedArray[ShapeT, np.dtype[np.void]],
    key: str | None = None,
    ignoremask: bool = True,
    *,
    return_index: Literal[True],
) -> tuple[np.ma.MaskedArray[ShapeT, np.dtype[np.void]], np.ndarray[ShapeT, np.dtype[np.int_]]]: ...

#
@overload
def join_by(
    key: str | Sequence[str],
    r1: npt.NDArray[np.void],
    r2: npt.NDArray[np.void],
    jointype: _JoinType = "inner",
    r1postfix: str = "1",
    r2postfix: str = "2",
    defaults: Mapping[str, object] | None = None,
    *,
    usemask: Literal[False],
    asrecarray: Literal[False] = False,
) -> np.ndarray[tuple[int], np.dtype[np.void]]: ...
@overload
def join_by(
    key: str | Sequence[str],
    r1: npt.NDArray[np.void],
    r2: npt.NDArray[np.void],
    jointype: _JoinType = "inner",
    r1postfix: str = "1",
    r2postfix: str = "2",
    defaults: Mapping[str, object] | None = None,
    *,
    usemask: Literal[False],
    asrecarray: Literal[True],
) -> np.recarray[tuple[int], np.dtype[np.void]]: ...
@overload
def join_by(
    key: str | Sequence[str],
    r1: npt.NDArray[np.void],
    r2: npt.NDArray[np.void],
    jointype: _JoinType = "inner",
    r1postfix: str = "1",
    r2postfix: str = "2",
    defaults: Mapping[str, object] | None = None,
    usemask: Literal[True] = True,
    asrecarray: Literal[False] = False,
) -> np.ma.MaskedArray[tuple[int], np.dtype[np.void]]: ...
@overload
def join_by(
    key: str | Sequence[str],
    r1: npt.NDArray[np.void],
    r2: npt.NDArray[np.void],
    jointype: _JoinType = "inner",
    r1postfix: str = "1",
    r2postfix: str = "2",
    defaults: Mapping[str, object] | None = None,
    usemask: Literal[True] = True,
    *,
    asrecarray: Literal[True],
) -> MaskedRecords[tuple[int], np.dtype[np.void]]: ...

#
def rec_join(
    key: str | Sequence[str],
    r1: npt.NDArray[np.void],
    r2: npt.NDArray[np.void],
    jointype: _JoinType = "inner",
    r1postfix: str = "1",
    r2postfix: str = "2",
    defaults: Mapping[str, object] | None = None,
) -> np.recarray[tuple[int], np.dtype[np.void]]: ...