"""Backend interface."""

from __future__ import annotations

from ..types import Array, Axis, Shape
from .device import Device


class Backend:
    """Backend protocol for array operations."""

    name: str
    pi: float
    e: float
    inf: float
    nan: float

    def name(self) -> str: ...
    def array(self, data: any, dtype: any = None, device: Device | None = None) -> Array: ...
    def convert(self, data: any, dtype: any = None, device: Device | None = None) -> Array: ...
    def randn(self, shape: Shape, dtype: any = None, device: Device | None = None) -> Array: ...
    def is_scalar(self, a: Array) -> bool: ...
    def add(self, a: Array, b: Array) -> Array: ...
    def subtract(self, a: Array, b: Array) -> Array: ...
    def multiply(self, a: Array, b: Array) -> Array: ...
    def divide(self, a: Array, b: Array) -> Array: ...
    def power(self, a: Array, b: Array) -> Array: ...
    def zeros(self, shape: Shape, dtype: any = None, device: Device | None = None) -> Array: ...
    def ones(self, shape: Shape, dtype: any = None, device: Device | None = None) -> Array: ...
    def full(
        self,
        shape: Shape,
        value: float,
        dtype: any = None,
        device: Device | None = None,
    ) -> Array: ...
    def arange(self, *args, dtype: any = None, device: Device | None = None) -> Array: ...
    def linspace(
        self,
        start: float,
        stop: float,
        num: int = 50,
        device: Device | None = None,
    ) -> Array: ...
    def eye(self, n: int, m: int | None = None, device: Device | None = None) -> Array: ...

    def random_key(self, seed: int) -> any: ...
    def random_normal(
        self,
        key: any,
        shape: Shape,
        dtype: any = None,
        device: Device | None = None,
    ) -> Array: ...
    def random_uniform(
        self,
        key: any,
        shape: Shape,
        low: float = 0.0,
        high: float = 1.0,
        dtype: any = None,
        device: Device | None = None,
    ) -> Array: ...
    def random_split(self, key: any, num: int = 2) -> list[any]: ...

    def reshape(self, a: Array, shape: Shape) -> Array: ...
    def transpose(self, a: Array, axes: list[int] | None = None) -> Array: ...
    def concat(self, arrays: list[Array], axis: int = 0) -> Array: ...
    def concatenate(self, arrays: list[Array], axis: int = 0) -> Array: ...
    def stack(self, arrays: list[Array], axis: int = 0) -> Array: ...
    def split(self, a: Array, indices: int | list[int], axis: int = 0) -> list[Array]: ...
    def squeeze(self, a: Array, axis: Axis = None) -> Array: ...
    def expand_dims(self, a: Array, axis: int) -> Array: ...
    def cumsum(self, a: Array, axis: int | None = None) -> Array: ...

    def sum(self, a: Array, axis: Axis = None, keepdims: bool = False) -> Array: ...
    def mean(self, a: Array, axis: Axis = None, keepdims: bool = False) -> Array: ...
    def std(self, a: Array, axis: Axis = None, keepdims: bool = False) -> Array: ...
    def var(self, a: Array, axis: Axis = None, keepdims: bool = False) -> Array: ...
    def min(self, a: Array, axis: Axis = None, keepdims: bool = False) -> Array: ...
    def max(self, a: Array, axis: Axis = None, keepdims: bool = False) -> Array: ...
    def argmin(self, a: Array, axis: Axis = None) -> Array: ...
    def argmax(self, a: Array, axis: Axis = None) -> Array: ...
    def prod(self, a: Array, axis: Axis = None, keepdims: bool = False) -> Array: ...

    def sub(self, a: Array, b: Array) -> Array: ...
    def mul(self, a: Array, b: Array) -> Array: ...
    def div(self, a: Array, b: Array) -> Array: ...
    def pow(self, a: Array, b: Array) -> Array: ...
    def matmul(self, a: Array, b: Array) -> Array: ...
    def dot(self, a: Array, b: Array) -> Array: ...

    def abs(self, a: Array) -> Array: ...
    def sign(self, a: Array) -> Array: ...
    def exp(self, a: Array) -> Array: ...
    def log(self, a: Array) -> Array: ...
    def log10(self, a: Array) -> Array: ...
    def sqrt(self, a: Array) -> Array: ...
    def sin(self, a: Array) -> Array: ...
    def cos(self, a: Array) -> Array: ...
    def tan(self, a: Array) -> Array: ...
    def tanh(self, a: Array) -> Array: ...
    def sinh(self, a: Array) -> Array: ...
    def cosh(self, a: Array) -> Array: ...

    def norm_cdf(self, x: Array) -> Array: ...
    def norm_pdf(self, x: Array) -> Array: ...
    def norm_ppf(self, q: Array) -> Array: ...
    def erf(self, x: Array) -> Array: ...
    def erfc(self, x: Array) -> Array: ...
    def gamma(self, x: Array) -> Array: ...
    def lgamma(self, x: Array) -> Array: ...

    def eq(self, a: Array, b: Array) -> Array: ...
    def ne(self, a: Array, b: Array) -> Array: ...
    def lt(self, a: Array, b: Array) -> Array: ...
    def le(self, a: Array, b: Array) -> Array: ...
    def gt(self, a: Array, b: Array) -> Array: ...
    def ge(self, a: Array, b: Array) -> Array: ...
    def isnan(self, a: Array) -> Array: ...
    def isinf(self, a: Array) -> Array: ...
    def isfinite(self, a: Array) -> Array: ...

    def any(self, a: Array, axis: Axis = None, keepdims: bool = False) -> Array: ...
    def all(self, a: Array, axis: Axis = None, keepdims: bool = False) -> Array: ...
    def logical_not(self, a: Array) -> Array: ...
    def logical_and(self, a: Array, b: Array) -> Array: ...
    def logical_or(self, a: Array, b: Array) -> Array: ...

    def where(self, cond: Array, x: Array, y: Array) -> Array: ...
    def clip(self, a: Array, min: float | None = None, max: float | None = None) -> Array: ...
    def maximum(self, a: Array, b: Array) -> Array: ...
    def minimum(self, a: Array, b: Array) -> Array: ...

    def cholesky(self, a: Array) -> Array: ...
    def solve(self, a: Array, b: Array) -> Array: ...
    def inv(self, a: Array) -> Array: ...
    def det(self, a: Array) -> Array: ...
    def eig(self, a: Array) -> tuple[Array, Array]: ...
    def svd(self, a: Array, full_matrices: bool = True) -> tuple[Array, Array, Array]: ...
    def qr(self, a: Array) -> tuple[Array, Array]: ...
    def norm(self, a: Array, ord: any = None, axis: Axis = None) -> Array: ...

    def shape(self, a: Array) -> Shape: ...
    def dtype(self, a: Array) -> any: ...
    def ndim(self, a: Array) -> int: ...
    def size(self, a: Array) -> int: ...
    def is_array(self, obj: any) -> bool: ...

    def to_numpy(self, a: Array) -> Array: ...
    def from_numpy(self, a: Array) -> Array: ...
    def item(self, a: Array) -> float | int | bool: ...
    def cast(self, a: Array, dtype: any) -> Array: ...
    def copy(self, a: Array) -> Array: ...

    def gather(self, a: Array, indices: Array, axis: int = 0) -> Array: ...
    def scatter(self, a: Array, indices: Array, values: Array, axis: int = 0) -> Array: ...
    def take(self, a: Array, indices: Array, axis: int | None = None) -> Array: ...
    def ensure_ndim(self, a: Array, ndim: int, message: str | None = None) -> Array: ...

    def memory_info(self, device: Device | None = None) -> dict[str, int]: ...
    def clear_cache(self, device: Device | None = None) -> None: ...


class AbstractBackend:
    """Base backend with common implementations."""

    name: str = "abstract"
    pi: float = 3.141592653589793
    e: float = 2.718281828459045
    inf: float = float("inf")
    nan: float = float("nan")

    def full(
        self,
        shape: Shape,
        value: float,
        dtype: any = None,
        device: Device | None = None,
    ) -> Array:
        """Array filled with value."""
        ones = self.ones(shape, dtype=dtype, device=device)
        return self.mul(ones, value)

    def eye(self, n: int, m: int | None = None, device: Device | None = None) -> Array:
        """Identity matrix."""
        m = m or n
        eye = self.zeros((n, m), device=device)
        # TODO: Implement proper eye matrix initialization
        # idx = self.arange(min(n, m), device=device)
        return eye

    def zeros_like(self, a: Array) -> Array:
        """Zeros with same shape/dtype."""
        return self.zeros(self.shape(a), dtype=self.dtype(a), device=self.device(a))

    def ones_like(self, a: Array) -> Array:
        """Ones with same shape/dtype."""
        return self.ones(self.shape(a), dtype=self.dtype(a), device=self.device(a))

    def ensure_ndim(self, a: Array, ndim: int, message: str | None = None) -> Array:
        """Ensure array has exactly ndim dimensions."""
        a = self.array(a)
        if self.ndim(a) != ndim:
            if self.ndim(a) < ndim:
                # Add leading dimensions
                for _ in range(ndim - self.ndim(a)):
                    a = self.expand_dims(a, 0)
            else:
                raise ValueError(message or f"Expected {ndim} dimensions, got {self.ndim(a)}")
        return a

    def convert(self, data: any, dtype: any = None, device: Device | None = None) -> Array:
        """Alias for array() for compatibility."""
        return self.array(data, dtype=dtype, device=device)

    def randn(self, shape: Shape, dtype: any = None, device: Device | None = None) -> Array:
        """Draw samples from standard normal distribution."""
        return self.random_normal(None, shape, dtype=dtype, device=device)

    def is_scalar(self, a: Array) -> bool:
        """Check if array is scalar."""
        return self.ndim(a) == 0

    def concatenate(self, arrays: list[Array], axis: int = 0) -> Array:
        """Alias for concat."""
        return self.concat(arrays, axis=axis)

    def cumsum(self, a: Array, axis: int | None = None) -> Array:
        """Compute cumulative sum."""
        # This will be implemented by child backends
        raise NotImplementedError("cumsum not implemented for this backend")

    # Aliases for convenience
    def add(self, a: Array, b: Array) -> Array: ...  # Handled by child or mixin
    def subtract(self, a: Array, b: Array) -> Array:
        return self.sub(a, b)

    def multiply(self, a: Array, b: Array) -> Array:
        return self.mul(a, b)

    def divide(self, a: Array, b: Array) -> Array:
        return self.div(a, b)

    def power(self, a: Array, b: Array) -> Array:
        return self.pow(a, b)

    def equal(self, a: Array, b: Array) -> Array:
        return self.eq(a, b)

    def not_equal(self, a: Array, b: Array) -> Array:
        return self.ne(a, b)

    def less(self, a: Array, b: Array) -> Array:
        return self.lt(a, b)

    def less_equal(self, a: Array, b: Array) -> Array:
        return self.le(a, b)

    def greater(self, a: Array, b: Array) -> Array:
        return self.gt(a, b)

    def greater_equal(self, a: Array, b: Array) -> Array:
        return self.ge(a, b)

    def copy(self, a: Array) -> Array:
        """Copy array."""
        return self.add(a, self.zeros_like(a))
