# Copyright 2025 hingebase

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:

# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.

# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.

# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

from typing import Any, TypeAlias, TypeVar, overload

import numpy as np
import numpy.typing as npt
from typing_extensions import TypeAliasType

from ._spec_array_object import array
from ._typing import Dtype, Shape

def abs(x: _ArrayT, /) -> _ArrayT: ...  # noqa: A001

@overload
def acos(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def acos(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def acos(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def acos(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

@overload
def acosh(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def acosh(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def acosh(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def acosh(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

def add(x1: array, x2: array, /) -> array: ...

@overload
def asin(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def asin(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def asin(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def asin(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

@overload
def asinh(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def asinh(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def asinh(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def asinh(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

@overload
def atan(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def atan(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def atan(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def atan(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

def atan2(x1: _RealArray, x2: _RealArray, /) -> _NDArray[np.floating[Any]]: ...

@overload
def atanh(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def atanh(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def atanh(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def atanh(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

def bitwise_and(
    x1: _BoolOrIntArray,
    x2: _BoolOrIntArray,
    /,
) -> _BoolOrIntArray: ...
def bitwise_invert(x: _BoolOrIntArrayT, /) -> _BoolOrIntArrayT: ...
def bitwise_left_shift(
    x1: _BoolOrIntArray,
    x2: _BoolOrIntArray,
    /,
) -> _NDArray[np.integer[Any]]: ...
def bitwise_or(
    x1: _BoolOrIntArray,
    x2: _BoolOrIntArray,
    /,
) -> _BoolOrIntArray: ...
def bitwise_right_shift(
    x1: _BoolOrIntArray,
    x2: _BoolOrIntArray,
    /,
) -> _NDArray[np.integer[Any]]: ...
def bitwise_xor(
    x1: _BoolOrIntArray,
    x2: _BoolOrIntArray,
    /,
) -> _BoolOrIntArray: ...

def ceil(x: _RealArrayT, /) -> _RealArrayT: ...
def conj(x: _NumericArrayT, /) -> _NumericArrayT: ...

@overload
def cos(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def cos(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def cos(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def cos(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

@overload
def cosh(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def cosh(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def cosh(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def cosh(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

def divide(x1: array, x2: array, /) -> _InexactArray: ...
def equal(x1: array, x2: array, /) -> _BoolArray: ...

@overload
def exp(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def exp(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def exp(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def exp(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

@overload
def expm1(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def expm1(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def expm1(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def expm1(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

def floor(x: _RealArrayT, /) -> _RealArrayT: ...
def floor_divide(
    x1: _RealArray,
    x2: _RealArray,
    /,
) -> _NDArray[np.integer[Any] | np.floating[Any]]: ...
def greater(x1: array, x2: array, /) -> _BoolArray: ...
def greater_equal(x1: array, x2: array, /) -> _BoolArray: ...

@overload
def imag(
    x: array[_ShapeT, Dtype[np.complexfloating[_BitT, _BitT]]],
    /,
) -> array[_ShapeT, Dtype[np.floating[_BitT]]]: ...
@overload
def imag(
    x: array[_ShapeT, Dtype[np.int32 | np.uint32 | np.float32]],
    /,
) -> _Float16Array[_ShapeT]: ...
@overload
def imag(
    x: array[_ShapeT, Dtype[np.int64 | np.uint64 | np.float64]],
    /,
) -> _Float32Array[_ShapeT]: ...

def isfinite(x: array[_ShapeT]) -> array[_ShapeT, Dtype[np.bool_]]: ...
def isinf(x: array[_ShapeT]) -> array[_ShapeT, Dtype[np.bool_]]: ...
def isnan(x: array[_ShapeT]) -> array[_ShapeT, Dtype[np.bool_]]: ...
def less(x1: array, x2: array, /) -> _BoolArray: ...
def less_equal(x1: array, x2: array, /) -> _BoolArray: ...

@overload
def log(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def log(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def log(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def log(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

@overload
def log1p(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def log1p(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def log1p(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def log1p(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

@overload
def log2(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def log2(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def log2(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def log2(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

@overload
def log10(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def log10(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def log10(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def log10(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

def logaddexp(x1: _RealArray, x2: _RealArray, /) -> _NDArray[np.floating[Any]]:
    ...
def logical_and(x1: array, x2: array, /) -> _BoolArray: ...
def logical_not(x: array[_ShapeT]) -> array[_ShapeT, Dtype[np.bool_]]: ...
def logical_or(x1: array, x2: array, /) -> _BoolArray: ...
def logical_xor(x1: array, x2: array, /) -> _BoolArray: ...
def multiply(x1: array, x2: array, /) -> array: ...
def negative(x: _NumericArrayT, /) -> _NumericArrayT: ...
def not_equal(x1: array, x2: array, /) -> _BoolArray: ...
def positive(x: _NumericArrayT, /) -> _NumericArrayT: ...
def pow(x1: array, x2: array, /) -> _NumericArray: ...  # noqa: A001

@overload
def real(
    x: array[_ShapeT, Dtype[np.complexfloating[_BitT, _BitT]]],
    /,
) -> array[_ShapeT, Dtype[np.floating[_BitT]]]: ...
@overload
def real(
    x: array[_ShapeT, Dtype[np.int32 | np.uint32 | np.float32]],
    /,
) -> _Float16Array[_ShapeT]: ...
@overload
def real(
    x: array[_ShapeT, Dtype[np.int64 | np.uint64 | np.float64]],
    /,
) -> _Float32Array[_ShapeT]: ...

def remainder(
    x1: _RealArray,
    x2: _RealArray,
    /,
) -> _NDArray[np.integer[Any] | np.floating[Any]]: ...

@overload
def round(x: _InexactArrayT, /) -> _InexactArrayT: ...  # noqa: A001
@overload
def round(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...  # noqa: A001
@overload
def round(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...  # noqa: A001
@overload
def round(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...  # noqa: A001

def sign(x: _NumericArrayT, /) -> _NumericArrayT: ...

@overload
def sin(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def sin(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def sin(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def sin(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

@overload
def sinh(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def sinh(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def sinh(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def sinh(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

@overload
def square(x: _NumericArrayT, /) -> _NumericArrayT: ...
@overload
def square(
    x: array[_ShapeT, Dtype[np.bool_]],
    /,
) -> array[_ShapeT, Dtype[np.int8]]: ...

@overload
def sqrt(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def sqrt(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def sqrt(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def sqrt(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

@overload
def subtract(x1: _NumericArray, x2: array, /) -> _NumericArray: ...
@overload
def subtract(x1: array, x2: _NumericArray, /) -> _NumericArray: ...

@overload
def tan(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def tan(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def tan(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def tan(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

@overload
def tanh(x: _InexactArrayT, /) -> _InexactArrayT: ...
@overload
def tanh(x: _ToFloat16Array[_ShapeT], /) -> _Float16Array[_ShapeT]: ...
@overload
def tanh(x: _ToFloat32Array[_ShapeT], /) -> _Float32Array[_ShapeT]: ...
@overload
def tanh(x: _ToFloat64Array[_ShapeT], /) -> _Float64Array[_ShapeT]: ...

def trunc(x: _RealArrayT, /) -> _RealArrayT: ...

_ArrayT = TypeVar("_ArrayT", bound=array)
_BitT = TypeVar("_BitT", bound=npt.NBitBase)
_SCT = TypeVar("_SCT", bound=np.bool_ | np.number[Any])
_NDArray = TypeAliasType(
    "_NDArray",
    array[Any, Dtype[_SCT]],
    type_params=(_SCT,),
)
_BoolArray: TypeAlias = _NDArray[np.bool_]
_BoolOrIntArray: TypeAlias = _NDArray[np.bool_ | np.integer[Any]]
_BoolOrIntArrayT = TypeVar("_BoolOrIntArrayT", bound=_BoolOrIntArray)
_ShapeT = TypeVar("_ShapeT", bound=Shape)
_Float16Array: TypeAlias = array[_ShapeT, Dtype[np.float16]]
_Float32Array: TypeAlias = array[_ShapeT, Dtype[np.float32]]
_Float64Array: TypeAlias = array[_ShapeT, Dtype[np.float64]]
_InexactArray: TypeAlias = _NDArray[np.inexact[Any]]
_InexactArrayT = TypeVar("_InexactArrayT", bound=_InexactArray)
_NumericArray: TypeAlias = _NDArray[np.number[Any]]
_NumericArrayT = TypeVar("_NumericArrayT", bound=_NumericArray)
_RealArray: TypeAlias = _NDArray[np.bool_ | np.integer[Any] | np.floating[Any]]
_RealArrayT = TypeVar("_RealArrayT", bound=_RealArray)
_ToFloat16Array: TypeAlias = array[
    _ShapeT,
    Dtype[np.int8 | np.uint8 | np.bool_],
]
_ToFloat32Array: TypeAlias = array[_ShapeT, Dtype[np.int16 | np.uint16]]
_ToFloat64Array: TypeAlias = array[
    _ShapeT,
    Dtype[np.int32 | np.uint32 | np.int64 | np.uint64],
]
