"""Type stubs for the fastdual._fastdual C extension."""

from typing import Any, overload

import numpy as np
import numpy.typing as npt


class Dual:
    """Dual number for forward-mode automatic differentiation."""

    val: float
    var_id: int
    grad: dict[int, float]
    real: "Dual"
    imag: "Dual"

    def __init__(self, value: float, seed: bool = True) -> None: ...

    # Arithmetic
    def __add__(self, other: "Dual | float | int") -> "Dual": ...
    def __radd__(self, other: float | int) -> "Dual": ...
    def __sub__(self, other: "Dual | float | int") -> "Dual": ...
    def __rsub__(self, other: float | int) -> "Dual": ...
    def __mul__(self, other: "Dual | float | int") -> "Dual": ...
    def __rmul__(self, other: float | int) -> "Dual": ...
    def __truediv__(self, other: "Dual | float | int") -> "Dual": ...
    def __rtruediv__(self, other: float | int) -> "Dual": ...
    def __floordiv__(self, other: "Dual | float | int") -> "Dual": ...
    def __rfloordiv__(self, other: float | int) -> "Dual": ...
    def __mod__(self, other: "Dual | float | int") -> "Dual": ...
    def __rmod__(self, other: float | int) -> "Dual": ...
    def __divmod__(self, other: "Dual | float | int") -> tuple["Dual", "Dual"]: ...
    def __rdivmod__(self, other: float | int) -> tuple["Dual", "Dual"]: ...
    def __pow__(self, other: "Dual | float | int", mod: None = None) -> "Dual": ...
    def __rpow__(self, other: float | int) -> "Dual": ...
    def __neg__(self) -> "Dual": ...
    def __pos__(self) -> "Dual": ...
    def __abs__(self) -> "Dual": ...
    def __float__(self) -> float: ...
    def __int__(self) -> int: ...
    def __bool__(self) -> bool: ...
    def __complex__(self) -> complex: ...

    # Comparison
    def __lt__(self, other: "Dual | float | int") -> bool: ...
    def __le__(self, other: "Dual | float | int") -> bool: ...
    def __eq__(self, other: object) -> bool: ...
    def __ne__(self, other: object) -> bool: ...
    def __gt__(self, other: "Dual | float | int") -> bool: ...
    def __ge__(self, other: "Dual | float | int") -> bool: ...
    def __hash__(self) -> int: ...

    # Formatting
    def __repr__(self) -> str: ...
    def __format__(self, spec: str) -> str: ...
    def __round__(self, ndigits: int | None = None) -> "Dual": ...
    def __floor__(self) -> "Dual": ...
    def __ceil__(self) -> "Dual": ...

    # Copy/pickle
    def __copy__(self) -> "Dual": ...
    def __deepcopy__(self, memo: dict[int, Any]) -> "Dual": ...
    def __reduce__(self) -> tuple[Any, ...]: ...

    # NumPy
    def __array_ufunc__(self, ufunc: Any, method: str, *inputs: Any, **kwargs: Any) -> Any: ...

    # Derivatives
    def der(self, wrt: "Dual") -> float: ...

    # Transcendentals
    def sin(self) -> "Dual": ...
    def cos(self) -> "Dual": ...
    def tan(self) -> "Dual": ...
    def exp(self) -> "Dual": ...
    def log(self) -> "Dual": ...
    def log2(self) -> "Dual": ...
    def log10(self) -> "Dual": ...
    def sqrt(self) -> "Dual": ...
    def arcsin(self) -> "Dual": ...
    def arccos(self) -> "Dual": ...
    def arctan(self) -> "Dual": ...
    def sinh(self) -> "Dual": ...
    def cosh(self) -> "Dual": ...
    def tanh(self) -> "Dual": ...
    def arcsinh(self) -> "Dual": ...
    def arccosh(self) -> "Dual": ...
    def arctanh(self) -> "Dual": ...
    def exp2(self) -> "Dual": ...
    def log1p(self) -> "Dual": ...
    def expm1(self) -> "Dual": ...
    def square(self) -> "Dual": ...
    def cbrt(self) -> "Dual": ...
    def sign(self) -> "Dual": ...
    def conjugate(self) -> "Dual": ...
    def fabs(self) -> "Dual": ...
    def floor(self) -> "Dual": ...
    def ceil(self) -> "Dual": ...


def reset() -> None:
    """Reset the global variable ID counter."""
    ...

def seed_array(values: list[float]) -> list[Dual]:
    """Create independent seed Duals from a list of floats (C-level)."""
    ...

def val_array(duals: list[Dual | float]) -> list[float]:
    """Extract primal values from a list of Duals."""
    ...

def der_array(results: list[Dual | float], wrt: Dual) -> list[float]:
    """Vectorized partial derivative extraction."""
    ...

def jac_matrix(
    results: list[Dual | float], seeds: list[Dual]
) -> tuple[list[float], int, int]:
    """Compute Jacobian in C. Returns (flat_list, n_rows, n_cols)."""
    ...

def apply_unary_array(
    name: str, arr: npt.NDArray[Any]
) -> npt.NDArray[Any] | None:
    """Apply a unary ufunc to an object array of Duals."""
    ...

def apply_binary_array(
    name: str,
    lhs: npt.NDArray[Any] | Dual | float,
    rhs: npt.NDArray[Any] | Dual | float,
) -> npt.NDArray[Any] | None:
    """Apply a binary ufunc to arrays/scalars of Duals."""
    ...
