from __future__ import annotations

from dataclasses import dataclass
from typing import (
    Any,
    Callable,
    ClassVar,
    Iterable,
    Mapping,
    Protocol,
    runtime_checkable,
)

from ..core.exceptions import BackendUnavailableError as BackendUnavailableError

__all__ = [
    "MathematicsBackend",
    "BackendUnavailableError",
    "register_backend",
    "get_backend",
    "available_backends",
    "ensure_array",
    "ensure_numpy",
]

@runtime_checkable
class MathematicsBackend(Protocol):
    name: str
    supports_autodiff: bool
    def as_array(self, value: Any, *, dtype: Any | None = None) -> Any: ...
    def eig(self, matrix: Any) -> tuple[Any, Any]: ...
    def eigh(self, matrix: Any) -> tuple[Any, Any]: ...
    def matrix_exp(self, matrix: Any) -> Any: ...
    def norm(
        self, value: Any, *, ord: Any | None = None, axis: Any | None = None
    ) -> Any: ...
    def einsum(self, pattern: str, *operands: Any, **kwargs: Any) -> Any: ...
    def matmul(self, a: Any, b: Any) -> Any: ...
    def conjugate_transpose(self, matrix: Any) -> Any: ...
    def stack(self, arrays: Iterable[Any], *, axis: int = 0) -> Any: ...
    def to_numpy(self, value: Any) -> Any: ...
    def is_gpu_available(self) -> bool: ...
    def get_device_name(self) -> str: ...
    def get_backend_info(self) -> Mapping[str, Any]: ...

BackendFactory = Callable[[], MathematicsBackend]

@dataclass(slots=True)
class _NumpyBackend:
    _np: Any
    _scipy_linalg: Any | None
    name: ClassVar[str] = ...
    supports_autodiff: ClassVar[bool] = ...
    def as_array(self, value: Any, *, dtype: Any | None = None) -> Any: ...
    def eig(self, matrix: Any) -> tuple[Any, Any]: ...
    def eigh(self, matrix: Any) -> tuple[Any, Any]: ...
    def matrix_exp(self, matrix: Any) -> Any: ...
    def norm(
        self, value: Any, *, ord: Any | None = None, axis: Any | None = None
    ) -> Any: ...
    def einsum(self, pattern: str, *operands: Any, **kwargs: Any) -> Any: ...
    def matmul(self, a: Any, b: Any) -> Any: ...
    def conjugate_transpose(self, matrix: Any) -> Any: ...
    def stack(self, arrays: Iterable[Any], *, axis: int = 0) -> Any: ...
    def to_numpy(self, value: Any) -> Any: ...
    def is_gpu_available(self) -> bool: ...
    def get_device_name(self) -> str: ...
    def get_backend_info(self) -> Mapping[str, Any]: ...

@dataclass(slots=True)
class _JaxBackend:
    _jnp: Any
    _jax_linalg: Any
    _jax: Any
    name: ClassVar[str] = ...
    supports_autodiff: ClassVar[bool] = ...
    def as_array(self, value: Any, *, dtype: Any | None = None) -> Any: ...
    def eig(self, matrix: Any) -> tuple[Any, Any]: ...
    def eigh(self, matrix: Any) -> tuple[Any, Any]: ...
    def matrix_exp(self, matrix: Any) -> Any: ...
    def norm(
        self, value: Any, *, ord: Any | None = None, axis: Any | None = None
    ) -> Any: ...
    def einsum(self, pattern: str, *operands: Any, **kwargs: Any) -> Any: ...
    def matmul(self, a: Any, b: Any) -> Any: ...
    def conjugate_transpose(self, matrix: Any) -> Any: ...
    def stack(self, arrays: Iterable[Any], *, axis: int = 0) -> Any: ...
    def to_numpy(self, value: Any) -> Any: ...
    def is_gpu_available(self) -> bool: ...
    def get_device_name(self) -> str: ...
    def get_backend_info(self) -> Mapping[str, Any]: ...

@dataclass(slots=True)
class _TorchBackend:
    _torch: Any
    _torch_linalg: Any
    _device: Any
    _use_cuda: bool
    name: ClassVar[str] = ...
    supports_autodiff: ClassVar[bool] = ...
    def as_array(self, value: Any, *, dtype: Any | None = None) -> Any: ...
    def eig(self, matrix: Any) -> tuple[Any, Any]: ...
    def eigh(self, matrix: Any) -> tuple[Any, Any]: ...
    def matrix_exp(self, matrix: Any) -> Any: ...
    def norm(
        self, value: Any, *, ord: Any | None = None, axis: Any | None = None
    ) -> Any: ...
    def einsum(self, pattern: str, *operands: Any, **kwargs: Any) -> Any: ...
    def matmul(self, a: Any, b: Any) -> Any: ...
    def conjugate_transpose(self, matrix: Any) -> Any: ...
    def stack(self, arrays: Iterable[Any], *, axis: int = 0) -> Any: ...
    def to_numpy(self, value: Any) -> Any: ...
    def is_gpu_available(self) -> bool: ...
    def get_device_name(self) -> str: ...
    def get_backend_info(self) -> Mapping[str, Any]: ...
    def get_device_info(self) -> dict[str, Any]: ...
    def to_cuda(self, value: Any) -> Any: ...
    def to_cpu(self, value: Any) -> Any: ...

def ensure_array(
    value: Any, *, dtype: Any | None = None, backend: MathematicsBackend | None = None
) -> Any: ...
def ensure_numpy(value: Any, *, backend: MathematicsBackend | None = None) -> Any: ...
def register_backend(
    name: str,
    factory: BackendFactory,
    *,
    aliases: Iterable[str] | None = None,
    override: bool = False,
) -> None: ...
def get_backend(name: str | None = None) -> MathematicsBackend: ...
def available_backends() -> Mapping[str, BackendFactory]: ...