from collections import deque
from pathlib import Path
from typing import Any, assert_type

import numpy as np
import numpy.typing as npt
from numpy._typing import _AnyShape

type _Array1D[ScalarT: np.generic] = np.ndarray[tuple[int], np.dtype[ScalarT]]
type _Array2D[ScalarT: np.generic] = np.ndarray[tuple[int, int], np.dtype[ScalarT]]
type _Array3D[ScalarT: np.generic] = np.ndarray[tuple[int, int, int], np.dtype[ScalarT]]

class SubClass[ScalarT: np.generic](np.ndarray[_AnyShape, np.dtype[ScalarT]]): ...

class IntoSubClass[ScalarT: np.generic]:
    def __array__(self) -> SubClass[ScalarT]: ...

i8: np.int64

A: npt.NDArray[np.float64]
B: SubClass[np.float64]
C: list[int]
D: SubClass[np.float64 | np.int64]
E: IntoSubClass[np.float64 | np.int64]

_f32_0d: np.float32
_f32_1d: _Array1D[np.float32]
_f32_2d: _Array2D[np.float32]
_f32_3d: _Array3D[np.float32]

_py_b_1d: list[bool]
_py_b_2d: list[list[bool]]
_py_i_1d: list[int]
_py_i_2d: list[list[int]]
_py_f_1d: list[float]
_py_f_2d: list[list[float]]
_py_c_1d: list[complex]
_py_c_2d: list[list[complex]]

mixed_shape: tuple[int, np.int64]

def func(i: int, j: int, **kwargs: Any) -> SubClass[np.float64]: ...

assert_type(np.empty_like(A), npt.NDArray[np.float64])
assert_type(np.empty_like(B), SubClass[np.float64])
assert_type(np.empty_like([1, 1.0]), npt.NDArray[Any])
assert_type(np.empty_like(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.empty_like(A, dtype="c16"), npt.NDArray[Any])

assert_type(np.array(A), npt.NDArray[np.float64])
assert_type(np.array(B), npt.NDArray[np.float64])
assert_type(np.array([1, 1.0]), npt.NDArray[Any])
assert_type(np.array(deque([1, 2, 3])), npt.NDArray[Any])
assert_type(np.array(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.array(A, dtype="c16"), npt.NDArray[Any])
assert_type(np.array(A, like=A), npt.NDArray[np.float64])
assert_type(np.array(A, subok=True), npt.NDArray[np.float64])
assert_type(np.array(B, subok=True), SubClass[np.float64])
assert_type(np.array(B, subok=True, ndmin=0), SubClass[np.float64])
assert_type(np.array(B, subok=True, ndmin=1), SubClass[np.float64])
assert_type(np.array(D), npt.NDArray[np.float64 | np.int64])
assert_type(np.array(E, subok=True), SubClass[np.float64 | np.int64])
# https://github.com/numpy/numpy/issues/29245
assert_type(np.array([], dtype=np.bool), npt.NDArray[np.bool])

assert_type(np.zeros([1, 5, 6]), npt.NDArray[np.float64])
assert_type(np.zeros([1, 5, 6], dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.zeros([1, 5, 6], dtype="c16"), npt.NDArray[Any])
assert_type(np.zeros(mixed_shape), npt.NDArray[np.float64])

assert_type(np.empty([1, 5, 6]), npt.NDArray[np.float64])
assert_type(np.empty([1, 5, 6], dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.empty([1, 5, 6], dtype="c16"), npt.NDArray[Any])
assert_type(np.empty(mixed_shape), npt.NDArray[np.float64])

assert_type(np.concatenate(A), npt.NDArray[np.float64])
assert_type(np.concatenate([A, A]), npt.NDArray[Any])  # pyright correctly infers this as NDArray[float64]
assert_type(np.concatenate([[1], A]), npt.NDArray[Any])
assert_type(np.concatenate([[1], [1]]), npt.NDArray[Any])
assert_type(np.concatenate((A, A)), npt.NDArray[np.float64])
assert_type(np.concatenate(([1], [1])), npt.NDArray[Any])
assert_type(np.concatenate([1, 1.0]), npt.NDArray[Any])
assert_type(np.concatenate(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.concatenate(A, dtype="c16"), npt.NDArray[Any])
assert_type(np.concatenate([1, 1.0], out=A), npt.NDArray[np.float64])

assert_type(np.asarray(A), npt.NDArray[np.float64])
assert_type(np.asarray(B), npt.NDArray[np.float64])
assert_type(np.asarray(C), _Array1D[np.int_])
assert_type(np.asarray(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.asarray(A, dtype="c16"), npt.NDArray[Any])
assert_type(np.asarray(_f32_1d), _Array1D[np.float32])
assert_type(np.asarray(_f32_1d, dtype=np.float64), _Array1D[np.float64])
assert_type(np.asarray(_f32_1d, dtype="f8"), _Array1D[Any])
assert_type(np.asarray([]), _Array1D[np.float64])
assert_type(np.asarray([[]]), _Array2D[np.float64])
assert_type(np.asarray(_py_b_1d), _Array1D[np.bool_])
assert_type(np.asarray(_py_b_2d), _Array2D[np.bool_])
assert_type(np.asarray(_py_i_1d), _Array1D[np.int_])
assert_type(np.asarray(_py_i_2d), _Array2D[np.int_])
assert_type(np.asarray(_py_f_1d), _Array1D[np.float64])
assert_type(np.asarray(_py_f_2d), _Array2D[np.float64])
assert_type(np.asarray(_py_c_1d), _Array1D[np.complex128])
assert_type(np.asarray(_py_c_2d), _Array2D[np.complex128])
assert_type(np.asarray(_py_i_1d, dtype=np.float32), _Array1D[np.float32])
assert_type(np.asarray(_py_i_1d, dtype="f4"), _Array1D[Any])
assert_type(np.asarray(_py_i_2d, dtype=np.float32), _Array2D[np.float32])
assert_type(np.asarray(_py_i_2d, dtype="f4"), _Array2D[Any])

assert_type(np.asanyarray(A), npt.NDArray[np.float64])
assert_type(np.asanyarray(B), SubClass[np.float64])
assert_type(np.asanyarray([1, 1.0]), npt.NDArray[Any])
assert_type(np.asanyarray(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.asanyarray(A, dtype="c16"), npt.NDArray[Any])

assert_type(np.ascontiguousarray(A), npt.NDArray[np.float64])
assert_type(np.ascontiguousarray(B), npt.NDArray[np.float64])
assert_type(np.ascontiguousarray([1, 1.0]), npt.NDArray[Any])
assert_type(np.ascontiguousarray(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.ascontiguousarray(A, dtype="c16"), npt.NDArray[Any])
assert_type(np.ascontiguousarray(_f32_1d), _Array1D[np.float32])

assert_type(np.asfortranarray(A), npt.NDArray[np.float64])
assert_type(np.asfortranarray(B), npt.NDArray[np.float64])
assert_type(np.asfortranarray([1, 1.0]), npt.NDArray[Any])
assert_type(np.asfortranarray(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.asfortranarray(A, dtype="c16"), npt.NDArray[Any])
assert_type(np.asfortranarray(_f32_1d), _Array1D[np.float32])

assert_type(np.fromstring("1 1 1", sep=" "), _Array1D[np.float64])
assert_type(np.fromstring(b"1 1 1", sep=" "), _Array1D[np.float64])
assert_type(np.fromstring("1 1 1", dtype=np.int64, sep=" "), _Array1D[np.int64])
assert_type(np.fromstring(b"1 1 1", dtype=np.int64, sep=" "), _Array1D[np.int64])
assert_type(np.fromstring("1 1 1", dtype="c16", sep=" "), _Array1D[Any])
assert_type(np.fromstring(b"1 1 1", dtype="c16", sep=" "), _Array1D[Any])

assert_type(np.fromfile("test.txt", sep=" "), _Array1D[np.float64])
assert_type(np.fromfile("test.txt", dtype=np.int64, sep=" "), _Array1D[np.int64])
assert_type(np.fromfile("test.txt", dtype="c16", sep=" "), _Array1D[Any])
with open("test.txt") as f:
    assert_type(np.fromfile(f, sep=" "), _Array1D[np.float64])
    assert_type(np.fromfile(b"test.txt", sep=" "), _Array1D[np.float64])
    assert_type(np.fromfile(Path("test.txt"), sep=" "), _Array1D[np.float64])

assert_type(np.fromiter("12345", np.float32), _Array1D[np.float32])
assert_type(np.fromiter("12345", np.float64), _Array1D[np.float64])
assert_type(np.fromiter("12345", bool), _Array1D[np.bool])
assert_type(np.fromiter("12345", int), _Array1D[np.int_ | Any])
assert_type(np.fromiter("12345", float), _Array1D[np.float64 | Any])
assert_type(np.fromiter("12345", complex), _Array1D[np.complex128 | Any])
assert_type(np.fromiter("12345", None), _Array1D[np.float64])
assert_type(np.fromiter("12345", object), _Array1D[Any])

assert_type(np.frombuffer(A), _Array1D[np.float64])
assert_type(np.frombuffer(A, dtype=np.int64), _Array1D[np.int64])
assert_type(np.frombuffer(A, dtype="c16"), _Array1D[Any])

_x_bool: bool
_x_int: int
_x_float: float
_x_timedelta: np.timedelta64[int]
_x_datetime: np.datetime64[int]

assert_type(np.arange(False, True), _Array1D[np.int_])
assert_type(np.arange(10), _Array1D[np.int_])
assert_type(np.arange(0, 10, step=2), _Array1D[np.int_])
assert_type(np.arange(10.0), _Array1D[np.float64 | Any])
assert_type(np.arange(0, stop=10.0), _Array1D[np.float64 | Any])
assert_type(np.arange(_x_timedelta), _Array1D[np.timedelta64])
assert_type(np.arange(0, _x_timedelta), _Array1D[np.timedelta64])
assert_type(np.arange(_x_datetime, _x_datetime), _Array1D[np.datetime64])
assert_type(np.arange(10, dtype=np.float64), _Array1D[np.float64])
assert_type(np.arange(0, 10, step=2, dtype=np.int16), _Array1D[np.int16])
assert_type(np.arange(10, dtype=int), _Array1D[np.int_])
assert_type(np.arange(0, 10, dtype="f8"), _Array1D[Any])
# https://github.com/numpy/numpy/issues/30628
assert_type(np.arange("2025-12-20", "2025-12-23", dtype="datetime64[D]"), _Array1D[np.datetime64])

assert_type(np.require(A), npt.NDArray[np.float64])
assert_type(np.require(B), SubClass[np.float64])
assert_type(np.require(B, requirements=None), SubClass[np.float64])
assert_type(np.require(B, dtype=int), npt.NDArray[Any])
assert_type(np.require(B, requirements="E"), npt.NDArray[Any])
assert_type(np.require(B, requirements=["ENSUREARRAY"]), npt.NDArray[Any])
assert_type(np.require(B, requirements={"F", "E"}), npt.NDArray[Any])
assert_type(np.require(B, requirements=["C", "OWNDATA"]), SubClass[np.float64])
assert_type(np.require(B, requirements="W"), SubClass[np.float64])
assert_type(np.require(B, requirements="A"), SubClass[np.float64])
assert_type(np.require(C), npt.NDArray[Any])

assert_type(np.linspace(0, 10), _Array1D[np.float64])
assert_type(np.linspace(0, 10j), _Array1D[np.complex128 | Any])
assert_type(np.linspace(0, 10, dtype=np.int64), _Array1D[np.int64])
assert_type(np.linspace(0, 10, dtype=int), npt.NDArray[Any])
assert_type(np.linspace(0, 10, retstep=True), tuple[_Array1D[np.float64], np.float64])
assert_type(np.linspace(0j, 10, retstep=True), tuple[_Array1D[np.complex128 | Any], np.complex128 | Any])
assert_type(np.linspace(0, 10, retstep=True, dtype=np.int64), tuple[_Array1D[np.int64], np.int64])
assert_type(np.linspace(0j, 10, retstep=True, dtype=int), tuple[npt.NDArray[Any], Any])

assert_type(np.logspace(0, 10), _Array1D[np.float64])
assert_type(np.logspace(0, 10j), _Array1D[np.complex128 | Any])
assert_type(np.logspace(0, 10, dtype=np.int64), _Array1D[np.int64])
assert_type(np.logspace(0, 10, dtype=int), npt.NDArray[Any])

assert_type(np.geomspace(0, 10), _Array1D[np.float64])
assert_type(np.geomspace(0, 10j), _Array1D[np.complex128 | Any])
assert_type(np.geomspace(0, 10, dtype=np.int64), _Array1D[np.int64])
assert_type(np.geomspace(0, 10, dtype=int), npt.NDArray[Any])

assert_type(np.zeros_like(A), npt.NDArray[np.float64])
assert_type(np.zeros_like(C), npt.NDArray[Any])
assert_type(np.zeros_like(A, dtype=float), npt.NDArray[Any])
assert_type(np.zeros_like(B), SubClass[np.float64])
assert_type(np.zeros_like(B, dtype=np.int64), npt.NDArray[np.int64])

assert_type(np.ones_like(A), npt.NDArray[np.float64])
assert_type(np.ones_like(C), npt.NDArray[Any])
assert_type(np.ones_like(A, dtype=float), npt.NDArray[Any])
assert_type(np.ones_like(B), SubClass[np.float64])
assert_type(np.ones_like(B, dtype=np.int64), npt.NDArray[np.int64])

assert_type(np.full_like(A, i8), npt.NDArray[np.float64])
assert_type(np.full_like(C, i8), npt.NDArray[Any])
assert_type(np.full_like(A, i8, dtype=int), npt.NDArray[Any])
assert_type(np.full_like(B, i8), SubClass[np.float64])
assert_type(np.full_like(B, i8, dtype=np.int64), npt.NDArray[np.int64])

_size: int
_shape_0d: tuple[()]
_shape_1d: tuple[int]
_shape_2d: tuple[int, int]
_shape_nd: tuple[int, ...]
_shape_like: list[int]

assert_type(np.ones(_shape_0d), np.ndarray[tuple[()], np.dtype[np.float64]])
assert_type(np.ones(_size), np.ndarray[tuple[int], np.dtype[np.float64]])
assert_type(np.ones(_shape_2d), np.ndarray[tuple[int, int], np.dtype[np.float64]])
assert_type(np.ones(_shape_nd), np.ndarray[tuple[int, ...], np.dtype[np.float64]])
assert_type(np.ones(_shape_1d, dtype=np.int64), np.ndarray[tuple[int], np.dtype[np.int64]])
assert_type(np.ones(_shape_like), npt.NDArray[np.float64])
assert_type(
    np.ones(_shape_like, dtype=np.dtypes.Int64DType()),
    np.ndarray[tuple[Any, ...], np.dtypes.Int64DType],
)
assert_type(np.ones(_shape_like, dtype=int), npt.NDArray[Any])
assert_type(np.ones(mixed_shape), npt.NDArray[np.float64])

assert_type(np.full(_size, i8), np.ndarray[tuple[int], np.dtype[np.int64]])
assert_type(np.full(_shape_2d, i8), np.ndarray[tuple[int, int], np.dtype[np.int64]])
assert_type(np.full(_shape_like, i8), npt.NDArray[np.int64])
assert_type(np.full(_shape_like, 42), npt.NDArray[Any])
assert_type(np.full(_size, i8, dtype=np.float64), np.ndarray[tuple[int], np.dtype[np.float64]])
assert_type(np.full(_size, i8, dtype=float), np.ndarray[tuple[int], np.dtype])
assert_type(np.full(_shape_like, 42, dtype=float), npt.NDArray[Any])
assert_type(np.full(_shape_0d, i8, dtype=object), np.ndarray[tuple[()], np.dtype])

assert_type(np.indices([1, 2, 3]), npt.NDArray[np.int_])
assert_type(np.indices([1, 2, 3], sparse=True), tuple[npt.NDArray[np.int_], ...])

assert_type(np.fromfunction(func, (3, 5)), SubClass[np.float64])

assert_type(np.identity(3), np.ndarray[tuple[int, int], np.dtype[np.float64]])
assert_type(np.identity(3, dtype=np.int8), np.ndarray[tuple[int, int], np.dtype[np.int8]])
assert_type(np.identity(3, dtype=bool), np.ndarray[tuple[int, int], np.dtype[np.bool]])
assert_type(np.identity(3, dtype="bool"), np.ndarray[tuple[int, int], np.dtype[np.bool]])
assert_type(np.identity(3, dtype="b1"), np.ndarray[tuple[int, int], np.dtype[np.bool]])
assert_type(np.identity(3, dtype="?"), np.ndarray[tuple[int, int], np.dtype[np.bool]])
assert_type(np.identity(3, dtype=int), np.ndarray[tuple[int, int], np.dtype[np.int_ | Any]])
assert_type(np.identity(3, dtype="int"), np.ndarray[tuple[int, int], np.dtype[np.int_ | Any]])
assert_type(np.identity(3, dtype="n"), np.ndarray[tuple[int, int], np.dtype[np.int_ | Any]])
assert_type(np.identity(3, dtype=float), np.ndarray[tuple[int, int], np.dtype[np.float64 | Any]])
assert_type(np.identity(3, dtype="float"), np.ndarray[tuple[int, int], np.dtype[np.float64 | Any]])
assert_type(np.identity(3, dtype="f8"), np.ndarray[tuple[int, int], np.dtype[np.float64 | Any]])
assert_type(np.identity(3, dtype="d"), np.ndarray[tuple[int, int], np.dtype[np.float64 | Any]])
assert_type(np.identity(3, dtype=complex), np.ndarray[tuple[int, int], np.dtype[np.complex128 | Any]])
assert_type(np.identity(3, dtype="complex"), np.ndarray[tuple[int, int], np.dtype[np.complex128 | Any]])
assert_type(np.identity(3, dtype="c16"), np.ndarray[tuple[int, int], np.dtype[np.complex128 | Any]])
assert_type(np.identity(3, dtype="D"), np.ndarray[tuple[int, int], np.dtype[np.complex128 | Any]])

assert_type(np.atleast_1d(_f32_0d), _Array1D[np.float32])
assert_type(np.atleast_1d(_f32_1d), _Array1D[np.float32])
assert_type(np.atleast_1d(_f32_2d), _Array2D[np.float32])
assert_type(np.atleast_1d(_f32_3d), _Array3D[np.float32])
assert_type(np.atleast_1d(A), npt.NDArray[np.float64])
assert_type(np.atleast_1d(C), npt.NDArray[Any])
assert_type(np.atleast_1d(A, A), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]])
assert_type(np.atleast_1d(A, C), tuple[npt.NDArray[Any], npt.NDArray[Any]])
assert_type(np.atleast_1d(C, C), tuple[npt.NDArray[Any], npt.NDArray[Any]])
assert_type(np.atleast_1d(A, A, A), tuple[npt.NDArray[np.float64], ...])
assert_type(np.atleast_1d(C, C, C), tuple[npt.NDArray[Any], ...])

assert_type(np.atleast_2d(_f32_0d), _Array2D[np.float32])
assert_type(np.atleast_2d(_f32_1d), _Array2D[np.float32])
assert_type(np.atleast_2d(_f32_2d), _Array2D[np.float32])
assert_type(np.atleast_2d(_f32_3d), _Array3D[np.float32])
assert_type(np.atleast_2d(A), npt.NDArray[np.float64])
assert_type(np.atleast_2d(A, A), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]])
assert_type(np.atleast_2d(A, A, A), tuple[npt.NDArray[np.float64], ...])

assert_type(np.atleast_3d(_f32_0d), _Array3D[np.float32])
assert_type(np.atleast_3d(_f32_1d), _Array3D[np.float32])
assert_type(np.atleast_3d(_f32_2d), _Array3D[np.float32])
assert_type(np.atleast_3d(_f32_3d), _Array3D[np.float32])
assert_type(np.atleast_3d(A), npt.NDArray[np.float64])
assert_type(np.atleast_3d(A, A), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]])
assert_type(np.atleast_3d(A, A, A), tuple[npt.NDArray[np.float64], ...])

assert_type(np.vstack([A, A]), npt.NDArray[np.float64])
assert_type(np.vstack([A, A], dtype=np.float32), npt.NDArray[np.float32])
assert_type(np.vstack([A, C]), npt.NDArray[Any])
assert_type(np.vstack([C, C]), npt.NDArray[Any])

assert_type(np.hstack([A, A]), npt.NDArray[np.float64])
assert_type(np.hstack([A, A], dtype=np.float32), npt.NDArray[np.float32])

assert_type(np.stack([A, A]), npt.NDArray[np.float64])
assert_type(np.stack([A, A], dtype=np.float32), npt.NDArray[np.float32])
assert_type(np.stack([A, C]), npt.NDArray[Any])
assert_type(np.stack([C, C]), npt.NDArray[Any])
assert_type(np.stack([A, A], axis=0), npt.NDArray[np.float64])
assert_type(np.stack([A, A], out=B), SubClass[np.float64])

assert_type(np.block([[A, A], [A, A]]), npt.NDArray[Any])  # pyright correctly infers this as NDArray[float64]
assert_type(np.block(C), npt.NDArray[Any])

from collections.abc import Buffer

def create_array(obj: npt.ArrayLike) -> npt.NDArray[Any]: ...

buffer: Buffer
assert_type(create_array(buffer), npt.NDArray[Any])
