import sys
from collections.abc import Sequence, Iterator, Callable, Iterable
from typing import (
    Literal as L,
    Any,
    TypeVar,
    overload,
    Protocol,
    SupportsIndex,
    SupportsInt,
)

if sys.version_info >= (3, 10):
    from typing import TypeGuard
else:
    from typing_extensions import TypeGuard

from numpy import (
    vectorize as vectorize,
    ufunc,
    generic,
    floating,
    complexfloating,
    intp,
    float64,
    complex128,
    timedelta64,
    datetime64,
    object_,
    _OrderKACF,
)

from numpy._typing import (
    NDArray,
    ArrayLike,
    DTypeLike,
    _ShapeLike,
    _ScalarLike_co,
    _DTypeLike,
    _ArrayLike,
    _ArrayLikeInt_co,
    _ArrayLikeFloat_co,
    _ArrayLikeComplex_co,
    _ArrayLikeTD64_co,
    _ArrayLikeDT64_co,
    _ArrayLikeObject_co,
    _FloatLike_co,
    _ComplexLike_co,
)

from numpy.core.function_base import (
    add_newdoc as add_newdoc,
)

from numpy.core.multiarray import (
    add_docstring as add_docstring,
    bincount as bincount,
)

from numpy.core.umath import _add_newdoc_ufunc

_T = TypeVar("_T")
_T_co = TypeVar("_T_co", covariant=True)
_SCT = TypeVar("_SCT", bound=generic)
_ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])

_2Tuple = tuple[_T, _T]

class _TrimZerosSequence(Protocol[_T_co]):
    def __len__(self) -> int: ...
    def __getitem__(self, key: slice, /) -> _T_co: ...
    def __iter__(self) -> Iterator[Any]: ...

class _SupportsWriteFlush(Protocol):
    def write(self, s: str, /) -> object: ...
    def flush(self) -> object: ...

__all__: list[str]

# NOTE: This is in reality a re-export of `np.core.umath._add_newdoc_ufunc`
def add_newdoc_ufunc(ufunc: ufunc, new_docstring: str, /) -> None: ...

@overload
def rot90(
    m: _ArrayLike[_SCT],
    k: int = ...,
    axes: tuple[int, int] = ...,
) -> NDArray[_SCT]: ...
@overload
def rot90(
    m: ArrayLike,
    k: int = ...,
    axes: tuple[int, int] = ...,
) -> NDArray[Any]: ...

@overload
def flip(m: _SCT, axis: None = ...) -> _SCT: ...
@overload
def flip(m: _ScalarLike_co, axis: None = ...) -> Any: ...
@overload
def flip(m: _ArrayLike[_SCT], axis: None | _ShapeLike = ...) -> NDArray[_SCT]: ...
@overload
def flip(m: ArrayLike, axis: None | _ShapeLike = ...) -> NDArray[Any]: ...

def iterable(y: object) -> TypeGuard[Iterable[Any]]: ...

@overload
def average(
    a: _ArrayLikeFloat_co,
    axis: None = ...,
    weights: None | _ArrayLikeFloat_co= ...,
    returned: L[False] = ...,
    keepdims: L[False] = ...,
) -> floating[Any]: ...
@overload
def average(
    a: _ArrayLikeComplex_co,
    axis: None = ...,
    weights: None | _ArrayLikeComplex_co = ...,
    returned: L[False] = ...,
    keepdims: L[False] = ...,
) -> complexfloating[Any, Any]: ...
@overload
def average(
    a: _ArrayLikeObject_co,
    axis: None = ...,
    weights: None | Any = ...,
    returned: L[False] = ...,
    keepdims: L[False] = ...,
) -> Any: ...
@overload
def average(
    a: _ArrayLikeFloat_co,
    axis: None = ...,
    weights: None | _ArrayLikeFloat_co= ...,
    returned: L[True] = ...,
    keepdims: L[False] = ...,
) -> _2Tuple[floating[Any]]: ...
@overload
def average(
    a: _ArrayLikeComplex_co,
    axis: None = ...,
    weights: None | _ArrayLikeComplex_co = ...,
    returned: L[True] = ...,
    keepdims: L[False] = ...,
) -> _2Tuple[complexfloating[Any, Any]]: ...
@overload
def average(
    a: _ArrayLikeObject_co,
    axis: None = ...,
    weights: None | Any = ...,
    returned: L[True] = ...,
    keepdims: L[False] = ...,
) -> _2Tuple[Any]: ...
@overload
def average(
    a: _ArrayLikeComplex_co | _ArrayLikeObject_co,
    axis: None | _ShapeLike = ...,
    weights: None | Any = ...,
    returned: L[False] = ...,
    keepdims: bool = ...,
) -> Any: ...
@overload
def average(
    a: _ArrayLikeComplex_co | _ArrayLikeObject_co,
    axis: None | _ShapeLike = ...,
    weights: None | Any = ...,
    returned: L[True] = ...,
    keepdims: bool = ...,
) -> _2Tuple[Any]: ...

@overload
def asarray_chkfinite(
    a: _ArrayLike[_SCT],
    dtype: None = ...,
    order: _OrderKACF = ...,
) -> NDArray[_SCT]: ...
@overload
def asarray_chkfinite(
    a: object,
    dtype: None = ...,
    order: _OrderKACF = ...,
) -> NDArray[Any]: ...
@overload
def asarray_chkfinite(
    a: Any,
    dtype: _DTypeLike[_SCT],
    order: _OrderKACF = ...,
) -> NDArray[_SCT]: ...
@overload
def asarray_chkfinite(
    a: Any,
    dtype: DTypeLike,
    order: _OrderKACF = ...,
) -> NDArray[Any]: ...

# TODO: Use PEP 612 `ParamSpec` once mypy supports `Concatenate`
# xref python/mypy#8645
@overload
def piecewise(
    x: _ArrayLike[_SCT],
    condlist: ArrayLike,
    funclist: Sequence[Any | Callable[..., Any]],
    *args: Any,
    **kw: Any,
) -> NDArray[_SCT]: ...
@overload
def piecewise(
    x: ArrayLike,
    condlist: ArrayLike,
    funclist: Sequence[Any | Callable[..., Any]],
    *args: Any,
    **kw: Any,
) -> NDArray[Any]: ...

def select(
    condlist: Sequence[ArrayLike],
    choicelist: Sequence[ArrayLike],
    default: ArrayLike = ...,
) -> NDArray[Any]: ...

@overload
def copy(
    a: _ArrayType,
    order: _OrderKACF,
    subok: L[True],
) -> _ArrayType: ...
@overload
def copy(
    a: _ArrayType,
    order: _OrderKACF = ...,
    *,
    subok: L[True],
) -> _ArrayType: ...
@overload
def copy(
    a: _ArrayLike[_SCT],
    order: _OrderKACF = ...,
    subok: L[False] = ...,
) -> NDArray[_SCT]: ...
@overload
def copy(
    a: ArrayLike,
    order: _OrderKACF = ...,
    subok: L[False] = ...,
) -> NDArray[Any]: ...

def gradient(
    f: ArrayLike,
    *varargs: ArrayLike,
    axis: None | _ShapeLike = ...,
    edge_order: L[1, 2] = ...,
) -> Any: ...

@overload
def diff(
    a: _T,
    n: L[0],
    axis: SupportsIndex = ...,
    prepend: ArrayLike = ...,
    append: ArrayLike = ...,
) -> _T: ...
@overload
def diff(
    a: ArrayLike,
    n: int = ...,
    axis: SupportsIndex = ...,
    prepend: ArrayLike = ...,
    append: ArrayLike = ...,
) -> NDArray[Any]: ...

@overload
def interp(
    x: _ArrayLikeFloat_co,
    xp: _ArrayLikeFloat_co,
    fp: _ArrayLikeFloat_co,
    left: None | _FloatLike_co = ...,
    right: None | _FloatLike_co = ...,
    period: None | _FloatLike_co = ...,
) -> NDArray[float64]: ...
@overload
def interp(
    x: _ArrayLikeFloat_co,
    xp: _ArrayLikeFloat_co,
    fp: _ArrayLikeComplex_co,
    left: None | _ComplexLike_co = ...,
    right: None | _ComplexLike_co = ...,
    period: None | _FloatLike_co = ...,
) -> NDArray[complex128]: ...

@overload
def angle(z: _ComplexLike_co, deg: bool = ...) -> floating[Any]: ...
@overload
def angle(z: object_, deg: bool = ...) -> Any: ...
@overload
def angle(z: _ArrayLikeComplex_co, deg: bool = ...) -> NDArray[floating[Any]]: ...
@overload
def angle(z: _ArrayLikeObject_co, deg: bool = ...) -> NDArray[object_]: ...

@overload
def unwrap(
    p: _ArrayLikeFloat_co,
    discont: None | float = ...,
    axis: int = ...,
    *,
    period: float = ...,
) -> NDArray[floating[Any]]: ...
@overload
def unwrap(
    p: _ArrayLikeObject_co,
    discont: None | float = ...,
    axis: int = ...,
    *,
    period: float = ...,
) -> NDArray[object_]: ...

def sort_complex(a: ArrayLike) -> NDArray[complexfloating[Any, Any]]: ...

def trim_zeros(
    filt: _TrimZerosSequence[_T],
    trim: L["f", "b", "fb", "bf"] = ...,
) -> _T: ...

@overload
def extract(condition: ArrayLike, arr: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
@overload
def extract(condition: ArrayLike, arr: ArrayLike) -> NDArray[Any]: ...

def place(arr: NDArray[Any], mask: ArrayLike, vals: Any) -> None: ...

def disp(
    mesg: object,
    device: None | _SupportsWriteFlush = ...,
    linefeed: bool = ...,
) -> None: ...

@overload
def cov(
    m: _ArrayLikeFloat_co,
    y: None | _ArrayLikeFloat_co = ...,
    rowvar: bool = ...,
    bias: bool = ...,
    ddof: None | SupportsIndex | SupportsInt = ...,
    fweights: None | ArrayLike = ...,
    aweights: None | ArrayLike = ...,
    *,
    dtype: None = ...,
) -> NDArray[floating[Any]]: ...
@overload
def cov(
    m: _ArrayLikeComplex_co,
    y: None | _ArrayLikeComplex_co = ...,
    rowvar: bool = ...,
    bias: bool = ...,
    ddof: None | SupportsIndex | SupportsInt = ...,
    fweights: None | ArrayLike = ...,
    aweights: None | ArrayLike = ...,
    *,
    dtype: None = ...,
) -> NDArray[complexfloating[Any, Any]]: ...
@overload
def cov(
    m: _ArrayLikeComplex_co,
    y: None | _ArrayLikeComplex_co = ...,
    rowvar: bool = ...,
    bias: bool = ...,
    ddof: None | SupportsIndex | SupportsInt = ...,
    fweights: None | ArrayLike = ...,
    aweights: None | ArrayLike = ...,
    *,
    dtype: _DTypeLike[_SCT],
) -> NDArray[_SCT]: ...
@overload
def cov(
    m: _ArrayLikeComplex_co,
    y: None | _ArrayLikeComplex_co = ...,
    rowvar: bool = ...,
    bias: bool = ...,
    ddof: None | SupportsIndex | SupportsInt = ...,
    fweights: None | ArrayLike = ...,
    aweights: None | ArrayLike = ...,
    *,
    dtype: DTypeLike,
) -> NDArray[Any]: ...

# NOTE `bias` and `ddof` have been deprecated
@overload
def corrcoef(
    m: _ArrayLikeFloat_co,
    y: None | _ArrayLikeFloat_co = ...,
    rowvar: bool = ...,
    *,
    dtype: None = ...,
) -> NDArray[floating[Any]]: ...
@overload
def corrcoef(
    m: _ArrayLikeComplex_co,
    y: None | _ArrayLikeComplex_co = ...,
    rowvar: bool = ...,
    *,
    dtype: None = ...,
) -> NDArray[complexfloating[Any, Any]]: ...
@overload
def corrcoef(
    m: _ArrayLikeComplex_co,
    y: None | _ArrayLikeComplex_co = ...,
    rowvar: bool = ...,
    *,
    dtype: _DTypeLike[_SCT],
) -> NDArray[_SCT]: ...
@overload
def corrcoef(
    m: _ArrayLikeComplex_co,
    y: None | _ArrayLikeComplex_co = ...,
    rowvar: bool = ...,
    *,
    dtype: DTypeLike,
) -> NDArray[Any]: ...

def blackman(M: _FloatLike_co) -> NDArray[floating[Any]]: ...

def bartlett(M: _FloatLike_co) -> NDArray[floating[Any]]: ...

def hanning(M: _FloatLike_co) -> NDArray[floating[Any]]: ...

def hamming(M: _FloatLike_co) -> NDArray[floating[Any]]: ...

def i0(x: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ...

def kaiser(
    M: _FloatLike_co,
    beta: _FloatLike_co,
) -> NDArray[floating[Any]]: ...

@overload
def sinc(x: _FloatLike_co) -> floating[Any]: ...
@overload
def sinc(x: _ComplexLike_co) -> complexfloating[Any, Any]: ...
@overload
def sinc(x: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ...
@overload
def sinc(x: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ...

@overload
def msort(a: _ArrayType) -> _ArrayType: ...
@overload
def msort(a: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
@overload
def msort(a: ArrayLike) -> NDArray[Any]: ...

@overload
def median(
    a: _ArrayLikeFloat_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    keepdims: L[False] = ...,
) -> floating[Any]: ...
@overload
def median(
    a: _ArrayLikeComplex_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    keepdims: L[False] = ...,
) -> complexfloating[Any, Any]: ...
@overload
def median(
    a: _ArrayLikeTD64_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    keepdims: L[False] = ...,
) -> timedelta64: ...
@overload
def median(
    a: _ArrayLikeObject_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    keepdims: L[False] = ...,
) -> Any: ...
@overload
def median(
    a: _ArrayLikeFloat_co | _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co,
    axis: None | _ShapeLike = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    keepdims: bool = ...,
) -> Any: ...
@overload
def median(
    a: _ArrayLikeFloat_co | _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co,
    axis: None | _ShapeLike = ...,
    out: _ArrayType = ...,
    overwrite_input: bool = ...,
    keepdims: bool = ...,
) -> _ArrayType: ...

_MethodKind = L[
    "inverted_cdf",
    "averaged_inverted_cdf",
    "closest_observation",
    "interpolated_inverted_cdf",
    "hazen",
    "weibull",
    "linear",
    "median_unbiased",
    "normal_unbiased",
    "lower",
    "higher",
    "midpoint",
    "nearest",
]

@overload
def percentile(
    a: _ArrayLikeFloat_co,
    q: _FloatLike_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    method: _MethodKind = ...,
    keepdims: L[False] = ...,
) -> floating[Any]: ...
@overload
def percentile(
    a: _ArrayLikeComplex_co,
    q: _FloatLike_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    method: _MethodKind = ...,
    keepdims: L[False] = ...,
) -> complexfloating[Any, Any]: ...
@overload
def percentile(
    a: _ArrayLikeTD64_co,
    q: _FloatLike_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    method: _MethodKind = ...,
    keepdims: L[False] = ...,
) -> timedelta64: ...
@overload
def percentile(
    a: _ArrayLikeDT64_co,
    q: _FloatLike_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    method: _MethodKind = ...,
    keepdims: L[False] = ...,
) -> datetime64: ...
@overload
def percentile(
    a: _ArrayLikeObject_co,
    q: _FloatLike_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    method: _MethodKind = ...,
    keepdims: L[False] = ...,
) -> Any: ...
@overload
def percentile(
    a: _ArrayLikeFloat_co,
    q: _ArrayLikeFloat_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    method: _MethodKind = ...,
    keepdims: L[False] = ...,
) -> NDArray[floating[Any]]: ...
@overload
def percentile(
    a: _ArrayLikeComplex_co,
    q: _ArrayLikeFloat_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    method: _MethodKind = ...,
    keepdims: L[False] = ...,
) -> NDArray[complexfloating[Any, Any]]: ...
@overload
def percentile(
    a: _ArrayLikeTD64_co,
    q: _ArrayLikeFloat_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    method: _MethodKind = ...,
    keepdims: L[False] = ...,
) -> NDArray[timedelta64]: ...
@overload
def percentile(
    a: _ArrayLikeDT64_co,
    q: _ArrayLikeFloat_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    method: _MethodKind = ...,
    keepdims: L[False] = ...,
) -> NDArray[datetime64]: ...
@overload
def percentile(
    a: _ArrayLikeObject_co,
    q: _ArrayLikeFloat_co,
    axis: None = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    method: _MethodKind = ...,
    keepdims: L[False] = ...,
) -> NDArray[object_]: ...
@overload
def percentile(
    a: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeTD64_co | _ArrayLikeObject_co,
    q: _ArrayLikeFloat_co,
    axis: None | _ShapeLike = ...,
    out: None = ...,
    overwrite_input: bool = ...,
    method: _MethodKind = ...,
    keepdims: bool = ...,
) -> Any: ...
@overload
def percentile(
    a: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeTD64_co | _ArrayLikeObject_co,
    q: _ArrayLikeFloat_co,
    axis: None | _ShapeLike = ...,
    out: _ArrayType = ...,
    overwrite_input: bool = ...,
    method: _MethodKind = ...,
    keepdims: bool = ...,
) -> _ArrayType: ...

# NOTE: Not an alias, but they do have identical signatures
# (that we can reuse)
quantile = percentile

# TODO: Returns a scalar for <= 1D array-likes; returns an ndarray otherwise
def trapz(
    y: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co,
    x: None | _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co = ...,
    dx: float = ...,
    axis: SupportsIndex = ...,
) -> Any: ...

def meshgrid(
    *xi: ArrayLike,
    copy: bool = ...,
    sparse: bool = ...,
    indexing: L["xy", "ij"] = ...,
) -> list[NDArray[Any]]: ...

@overload
def delete(
    arr: _ArrayLike[_SCT],
    obj: slice | _ArrayLikeInt_co,
    axis: None | SupportsIndex = ...,
) -> NDArray[_SCT]: ...
@overload
def delete(
    arr: ArrayLike,
    obj: slice | _ArrayLikeInt_co,
    axis: None | SupportsIndex = ...,
) -> NDArray[Any]: ...

@overload
def insert(
    arr: _ArrayLike[_SCT],
    obj: slice | _ArrayLikeInt_co,
    values: ArrayLike,
    axis: None | SupportsIndex = ...,
) -> NDArray[_SCT]: ...
@overload
def insert(
    arr: ArrayLike,
    obj: slice | _ArrayLikeInt_co,
    values: ArrayLike,
    axis: None | SupportsIndex = ...,
) -> NDArray[Any]: ...

def append(
    arr: ArrayLike,
    values: ArrayLike,
    axis: None | SupportsIndex = ...,
) -> NDArray[Any]: ...

@overload
def digitize(
    x: _FloatLike_co,
    bins: _ArrayLikeFloat_co,
    right: bool = ...,
) -> intp: ...
@overload
def digitize(
    x: _ArrayLikeFloat_co,
    bins: _ArrayLikeFloat_co,
    right: bool = ...,
) -> NDArray[intp]: ...