# This minimal stub allows pytype to do some basic type-checking of numpy. We
# should replace it with the official numpy stubs.

import sys
from typing import Any, Generic, Iterator, Mapping, Optional, Tuple, TypeVar

# Some type differences under python3
if sys.version_info[0] >= 3:
  long = int
  buffer_type = memoryview
else:
  buffer_type = buffer

_ShapeType = TypeVar("_ShapeType", bound=Any)
_DType_co = TypeVar("_DType_co", covariant=True, bound=dtype[Any])
_ArraySelf = TypeVar("_ArraySelf", bound=ndarray)

class ndarray(Generic[_ShapeType, _DType_co]):
  # https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.html
  def __new__(cls: type[_ArraySelf],
              shape,
              dtype=...,
              buffer=...,
              offset=...,
              strides=...,
              order=...) -> _ArraySelf: ...

  def __getattr__(self, name) -> Any: ...

  # We need to add all special methods used for protocols explicitly
  # (b/33779140). Everything else is optional.

  # https://docs.scipy.org/doc/numpy/reference/arrays.ndarray.html#array-attributes
  @property
  def flags(self) -> Any: ...
  @property
  def shape(self) -> Tuple[int, ...]: ...
  @property
  def strides(self) -> Tuple[int, ...]: ...
  @property
  def data(self) -> buffer_type: ...
  @property
  def size(self) -> int: ...
  @property
  def itemsize(self) -> int: ...
  @property
  def nbytes(self) -> int: ...
  @property
  def base(self) -> Optional[ndarray]: ...

  @property
  def T(self) -> ndarray: ...
  @property
  def real(self) -> ndarray: ...
  @property
  def imag(self) -> ndarray: ...
  @property
  def flat(self) -> Any: ...
  @property
  def ctypes(self) -> Any: ...
  @property
  def __array_interface__(self) -> dict: ...

  # Array methods are missing:
  # https://docs.scipy.org/doc/numpy/reference/arrays.ndarray.html#array-methods

  # https://docs.python.org/2/reference/datamodel.html
  def __len__(self) -> int: ...
  def __getitem__(self, key) -> Any: ...
  def __setitem__(self, key, value): ...
  def __iter__(self) -> Any: ...
  def __contains__(self, key) -> bool: ...
  def __reversed__(self) -> Iterator[Any]: ...

  # __getitem__ should suffice here, but pytype needs these (b/65407594)
  def __getslice__(self, i, j) -> Any: ...
  def __setslice__(self, i, j, value) -> Any: ...

  def __complex__(self) -> complex: ...
  def __int__(self) -> int: ...
  def __long__(self) -> long: ...
  def __float__(self) -> float: ...
  def __oct__(self) -> str: ...
  def __hex__(self) -> str: ...
  def __nonzero__(self) -> bool: ...

  def __index__(self) -> int: ...

  def __copy__(self, order = ...) -> ndarray: ...
  def __deepcopy__(self, order = ...) -> ndarray: ...

  # https://github.com/numpy/numpy/blob/v1.13.0/numpy/lib/mixins.py#L63-L181

  # We ought to add overloads (returning ndarray) for cases where other is known
  # not to define __array_priority__ or __array_ufunc__, such as for numbers or
  # other numpy arrays. Or even better, use protocols (once they work).

  def __lt__(self, other) -> ndarray: ...
  def __le__(self, other) -> ndarray: ...
  def __eq__(self, other): ...
  def __ne__(self, other): ...
  def __gt__(self, other) -> ndarray: ...
  def __ge__(self, other) -> ndarray: ...

  @overload
  def __add__(self, other: complex) -> ndarray: ...
  @overload
  def __add__(self, other): ...
  @overload
  def __radd__(self, other: complex) -> ndarray: ...
  @overload
  def __radd__(self, other): ...
  @overload
  def __iadd__(self, other: complex) -> ndarray: ...
  @overload
  def __iadd__(self, other): ...

  @overload
  def __sub__(self, other: complex) -> ndarray: ...
  @overload
  def __sub__(self, other): ...
  @overload
  def __rsub__(self, other: complex) -> ndarray: ...
  @overload
  def __rsub__(self, other): ...
  @overload
  def __isub__(self, other: complex) -> ndarray: ...
  @overload
  def __isub__(self, other): ...

  @overload
  def __mul__(self, other: complex) -> ndarray: ...
  @overload
  def __mul__(self, other): ...
  @overload
  def __rmul__(self, other: complex) -> ndarray: ...
  @overload
  def __rmul__(self, other): ...
  @overload
  def __imul__(self, other: complex) -> ndarray: ...
  @overload
  def __imul__(self, other): ...

  def __div__(self, other): ...
  def __rdiv__(self, other): ...
  def __idiv__(self, other): ...

  @overload
  def __truediv__(self, other: complex) -> ndarray: ...
  @overload
  def __truediv__(self, other): ...
  @overload
  def __rtruediv__(self, other: complex) -> ndarray: ...
  @overload
  def __rtruediv__(self, other): ...
  @overload
  def __itruediv__(self, other: complex) -> ndarray: ...
  @overload
  def __itruediv__(self, other): ...

  @overload
  def __floordiv__(self, other: float) -> ndarray: ...
  @overload
  def __floordiv__(self, other): ...
  @overload
  def __rfloordiv__(self, other: float) -> ndarray: ...
  @overload
  def __rfloordiv__(self, other): ...
  @overload
  def __ifloordiv__(self, other: float) -> ndarray: ...
  @overload
  def __ifloordiv__(self, other): ...

  @overload
  def __mod__(self, other: float) -> ndarray: ...
  @overload
  def __mod__(self, other): ...
  @overload
  def __rmod__(self, other: float) -> ndarray: ...
  @overload
  def __rmod__(self, other): ...
  @overload
  def __imod__(self, other: float) -> ndarray: ...
  @overload
  def __imod__(self, other): ...

  def __divmod__(self, other) -> tuple[ndarray, ndarray]: ...
  def __rdivmod__(self, other) -> tuple[ndarray, ndarray]: ...

  # NumPy's __pow__ doesn't handle a third argument
  @overload
  def __pow__(self, other: complex) -> ndarray: ...
  @overload
  def __pow__(self, other): ...
  @overload
  def __rpow__(self, other: complex) -> ndarray: ...
  @overload
  def __rpow__(self, other): ...
  @overload
  def __ipow__(self, other: complex) -> ndarray: ...
  @overload
  def __ipow__(self, other): ...

  @overload
  def __lshift__(self, other: int) -> ndarray: ...
  @overload
  def __lshift__(self, other): ...
  @overload
  def __rlshift__(self, other: int) -> ndarray: ...
  @overload
  def __rlshift__(self, other): ...
  @overload
  def __ilshift__(self, other: int) -> ndarray: ...
  @overload
  def __ilshift__(self, other): ...

  @overload
  def __rshift__(self, other: int) -> ndarray: ...
  @overload
  def __rshift__(self, other): ...
  @overload
  def __rrshift__(self, other: int) -> ndarray: ...
  @overload
  def __rrshift__(self, other): ...
  @overload
  def __irshift__(self, other: int) -> ndarray: ...
  @overload
  def __irshift__(self, other): ...

  @overload
  def __and__(self, other: int) -> ndarray: ...
  @overload
  def __and__(self, other): ...
  @overload
  def __rand__(self, other: int) -> ndarray: ...
  @overload
  def __rand__(self, other): ...
  @overload
  def __iand__(self, other: int) -> ndarray: ...
  @overload
  def __iand__(self, other): ...

  @overload
  def __xor__(self, other: int) -> ndarray: ...
  @overload
  def __xor__(self, other): ...
  @overload
  def __rxor__(self, other: int) -> ndarray: ...
  @overload
  def __rxor__(self, other): ...
  @overload
  def __ixor__(self, other: int) -> ndarray: ...
  @overload
  def __ixor__(self, other): ...

  @overload
  def __or__(self, other: int) -> ndarray: ...
  @overload
  def __or__(self, other): ...
  @overload
  def __ror__(self, other: int) -> ndarray: ...
  @overload
  def __ror__(self, other): ...
  @overload
  def __ior__(self, other: int) -> ndarray: ...
  @overload
  def __ior__(self, other): ...

  def __neg__(self) -> ndarray: ...
  def __pos__(self) -> ndarray: ...
  def __abs__(self) -> ndarray: ...
  def __invert__(self) -> ndarray: ...

  @overload
  def __matmul__(self, other: complex) -> ndarray: ...
  @overload
  def __matmul__(self, other): ...
  @overload
  def __rmatmul__(self, other: complex) -> ndarray: ...
  @overload
  def __rmatmul__(self, other): ...
  @overload
  def __imatmul__(self, other: complex) -> ndarray: ...
  @overload
  def __imatmul__(self, other): ...

def array(object: Any,
          dtype: Any = ...,
          copy: bool = ...,
          order: str = ...,
          subok: bool = ...,
          ndmin: int = ...) -> ndarray: ...

def asarray(a: Any,
            dtype: Any = ...,
            order: Optional[str] = ...) -> ndarray: ...

_DTypeScalar_co = TypeVar("_DTypeScalar_co", covariant=True)

class dtype(Generic[_DTypeScalar_co]):
    names: None | tuple[builtins.str, ...]

    def __init__(self, *args, **kwargs) -> None: ...

    if sys.version_info >= (3, 9):
        def __class_getitem__(self, item): ...

    def __getitem__(self, key) -> dtype[Any]: ...
    def __mul__(self, value): ...
    def __rmul__(self, value): ...
    def __gt__(self, other) -> bool: ...
    def __ge__(self, other) -> bool: ...
    def __lt__(self, other) -> bool: ...
    def __le__(self, other) -> bool: ...
    def __eq__(self, other) -> bool: ...
    def __ne__(self, other) -> bool: ...

    @property
    def alignment(self) -> int: ...
    @property
    def base(self) -> dtype[Any]: ...
    @property
    def byteorder(self) -> str: ...
    @property
    def char(self) -> str: ...
    @property
    def descr(self) -> list[tuple[str, str] | tuple[str, str, Any]]: ...
    @property
    def fields(
        self,
    ) -> None | Mapping[str, tuple[dtype[Any], int] | tuple[dtype[Any], int, Any]]: ...
    @property
    def flags(self) -> int: ...
    @property
    def hasobject(self) -> bool: ...
    @property
    def isbuiltin(self) -> int: ...
    @property
    def isnative(self) -> bool: ...
    @property
    def isalignedstruct(self) -> bool: ...
    @property
    def itemsize(self) -> int: ...
    @property
    def kind(self) -> str: ...
    @property
    def metadata(self) -> None | Mapping[str, Any]: ...
    @property
    def name(self) -> str: ...
    @property
    def num(self) -> int: ...
    @property
    def shape(self): ...
    @property
    def ndim(self) -> int: ...
    @property
    def subdtype(self) -> None | tuple[dtype[Any], Any]: ...
    def newbyteorder(self, __new_order = ...): ...
    @property
    def str(self) -> str: ...
    @property
    def type(self) -> type[_DTypeScalar_co]: ...

def __getattr__(name) -> Any: ...
