from __future__ import annotations

import copy
import operator
from typing import Callable

import numpy as np

from ..utils import ShapeError


r_tol: float = 1e-9
a_tol: float = 1e-12

# Backward-compatible names used in earlier revisions.
relative_tolerance: float = r_tol
absolute_tolerance: float = a_tol


def format_float(value: float) -> str:
    return np.format_float_positional(float(value), trim="-")


def format_array(data) -> str:
    return np.array2string(
        data,
        formatter={
            "float_kind": format_float,
        },
    )


_format_array = format_array


class ArrayBase:
    __slots__ = ["data"]

    def __init__(self, data: np.ndarray | list | tuple) -> None:
        self.data = np.asarray(data, dtype=np.float64)

        if self.__class__.__name__ == "Vector" and self.ndim != 1:
            raise ValueError("Vector must be 1D")

        if self.__class__.__name__ == "Matrix" and self.ndim != 2:
            raise ValueError("Matrix must be 2D")

        if self.__class__.__name__ == "Tensor" and self.ndim < 3:
            raise ValueError("Tensor must be 3D or higher")

    # ---- info ----

    @property
    def shape(self) -> tuple[int, ...]:
        return self.data.shape

    @property
    def ndim(self) -> int:
        return self.data.ndim

    @property
    def size(self) -> int:
        return self.data.size

    @property
    def dtype(self) -> np.dtype:
        return self.data.dtype

    # ---- representation ----

    def __repr__(self):
        return f"{self.__class__.__name__}({format_array(self.data)})"

    def __str__(self):
        return format_array(self.data)

    def __bool__(self):
        """Return True if the array is not empty."""
        return self.size > 0

    def __float__(self):
        """Return a Python float if the array is a single scalar."""
        if self.size != 1:
            raise ValueError(f"Cannot convert {self.__class__.__name__} with shape {self.shape} to float")
        return float(self.data)

    def __complex__(self):
        """Return a Python complex if the array is a single scalar."""
        if self.size != 1:
            raise ValueError(f"Cannot convert {self.__class__.__name__} with shape {self.shape} to complex")
        return complex(self.data)

    def __array__(self, dtype=None) -> np.ndarray:
        """Return the underlying NumPy array (for interoperability)."""
        if dtype is not None:
            return self.data.astype(dtype)
        return self.data

    def __format__(self, format_spec) -> str:
        """Format the array using the specified format specifier."""
        formatted = np.array2string(
            self.data,
            formatter={
                "float_kind": lambda x: format(float(x), format_spec),
                "complex_kind": lambda x: format(complex(x), format_spec),
            },
        )
        return f"{self.__class__.__name__}({formatted})"

    # ---- length and iteration ----

    def __len__(self) -> int:
        """Return the size of the first dimension."""
        return self.data.shape[0]

    def __iter__(self) -> iter:
        """Return an iterator over the first dimension."""
        return iter(self.data)

    def __contains__(self, item) -> bool:
        """Check if item is in the array."""
        if np.isscalar(item):
            return bool(np.any(np.isclose(self.data, item, rtol=r_tol, atol=a_tol)))
        return item in self.data

    def __setitem__(self, key, value) -> None:
        """Set item or slice in the array."""
        self.data[key] = value

    def __getitem__(self, key):
        """Get item or slice from the array."""
        result = self.data[key]
        # Return Python-native scalar (no trailing dot) for scalar indexing
        if np.isscalar(result) or result.ndim == 0:
            val = float(result)
            return int(val) if val.is_integer() else val
        return self.__class__(result)

    # ---- comparison operators ----

    def _require_same_type(self, other: object) -> bool:
        return isinstance(other, self.__class__)

    def _require_same_shape(self, other: ArrayBase) -> None:
        if self.shape != other.shape:
            raise ShapeError(f"Expected arrays of the same shape, but got {self.shape} and {other.shape}")

    def _require_nonzero(self, data: np.ndarray) -> None:
        if np.any(np.isclose(data, 0.0, rtol=r_tol, atol=a_tol)):
            raise ZeroDivisionError("Cannot divide by zero")

    def _binary_elementwise(self, other: object, op: Callable[[np.ndarray, np.ndarray], np.ndarray]):
        if not self._require_same_type(other):
            return NotImplemented
        self._require_same_shape(other)
        return self.__class__(op(self.data, other.data))

    def _binary_compare(self, other: object, op: Callable[[np.ndarray, np.ndarray], np.ndarray]):
        if not self._require_same_type(other):
            return NotImplemented
        return self.__class__(op(self.data, other.data))

    def _inplace_elementwise(self, other: object, op: Callable[[np.ndarray, np.ndarray], np.ndarray]):
        if not self._require_same_type(other):
            return NotImplemented
        self._require_same_shape(other)
        op(self.data, other.data)
        return self

    def __eq__(self, other) -> ArrayBase:
        """Element-wise equality comparison."""
        if not self._require_same_type(other):
            return NotImplemented
        return self.__class__(np.isclose(self.data, other.data, rtol=r_tol, atol=a_tol))

    def __ne__(self, other) -> ArrayBase:
        """Element-wise inequality comparison."""
        if not self._require_same_type(other):
            return NotImplemented
        return self.__class__(np.logical_not(np.isclose(self.data, other.data, rtol=r_tol, atol=a_tol)))

    def __lt__(self, other) -> ArrayBase:
        """Element-wise less than comparison."""
        return self._binary_compare(other, operator.lt)

    def __le__(self, other) -> ArrayBase:
        """Element-wise less than or equal comparison."""
        return self._binary_compare(other, operator.le)

    def __gt__(self, other) -> ArrayBase:
        """Element-wise greater than comparison."""
        return self._binary_compare(other, operator.gt)

    def __ge__(self, other) -> ArrayBase:
        """Element-wise greater than or equal comparison."""
        return self._binary_compare(other, operator.ge)

    # ---- unary operators ----

    def __neg__(self) -> ArrayBase:
        """Unary negation."""
        if not isinstance(self, self.__class__):
            return NotImplemented
        return self.__class__(-self.data)

    def __pos__(self) -> ArrayBase:
        """Unary positive."""
        if not isinstance(self, self.__class__):
            return NotImplemented
        return self.__class__(+self.data)

    def __abs__(self) -> ArrayBase:
        """Absolute value."""
        if not isinstance(self, self.__class__):
            return NotImplemented
        return self.__class__(np.abs(self.data))

    # ---- copy methods ----

    def __copy__(self) -> ArrayBase:
        """Shallow copy."""
        return self.__class__(self.data.copy())

    def __deepcopy__(self, memo) -> ArrayBase:
        """Deep copy."""
        return self.__class__(copy.deepcopy(self.data, memo))

    def copy(self) -> ArrayBase:
        """Return a copy of the array."""
        return self.__copy__()

    # ---- common array methods ----

    def any(self) -> bool:
        """Return True if any element is True."""
        return self.data.any()

    def all(self) -> bool:
        """Return True if all elements are True."""
        return self.data.all()

    def reshape(self, *shape) -> ArrayBase:
        """Return a new array with the specified shape."""
        return self.__class__(self.data.reshape(*shape))

    def flatten(self) -> ArrayBase:
        """Return a flattened 1-D array."""
        return self.__class__(self.data.flatten())

    def astype(self, dtype) -> ArrayBase:
        """Return array with specified dtype."""
        return self.__class__(self.data.astype(dtype))

    # ---- arithmetic operators ----

    def __add__(self, other: ArrayBase) -> ArrayBase:
        """Element-wise addition."""
        return self._binary_elementwise(other, operator.add)

    def __sub__(self, other: ArrayBase) -> ArrayBase:
        """Element-wise subtraction."""
        return self._binary_elementwise(other, operator.sub)

    def __mul__(self, other: ArrayBase) -> ArrayBase:
        """Element-wise multiplication."""
        return self._binary_elementwise(other, operator.mul)

    def __truediv__(self, other: ArrayBase) -> ArrayBase:
        """Element-wise division."""
        if not self._require_same_type(other):
            return NotImplemented
        self._require_same_shape(other)
        self._require_nonzero(other.data)
        return self.__class__(self.data / other.data)

    def __matmul__(self, other: ArrayBase) -> ArrayBase:
        """Matrix multiplication."""
        if not isinstance(other, ArrayBase):
            return NotImplemented
        if self.shape[-1] != other.shape[0]:
            raise ValueError(f"{self.__class__.__name__} shapes are not aligned for matrix multiplication")
        result = self.data @ other.data
        if np.isscalar(result) or getattr(result, "ndim", 0) == 0:
            return float(result)
        # If dimensionality drops (e.g. Matrix @ Vector -> 1-D), return the
        # lower-dimensional operand's class (Vector) rather than forcing Matrix.
        if result.ndim == other.data.ndim:
            return other.__class__(result)
        return self.__class__(result)

    # ---- right-hand arithmetic operators ----

    def __radd__(self, other: ArrayBase) -> ArrayBase:
        """Right-hand element-wise addition."""
        return self._binary_elementwise(other, lambda a, b: operator.add(b, a))

    def __rsub__(self, other: ArrayBase) -> ArrayBase:
        """Right-hand element-wise subtraction."""
        return self._binary_elementwise(other, lambda a, b: operator.sub(b, a))

    def __rmul__(self, other: ArrayBase) -> ArrayBase:
        """Right-hand element-wise multiplication."""
        return self._binary_elementwise(other, lambda a, b: operator.mul(b, a))

    def __rtruediv__(self, other: ArrayBase) -> ArrayBase:
        """Right-hand element-wise division."""
        if not self._require_same_type(other):
            return NotImplemented
        self._require_same_shape(other)
        self._require_nonzero(self.data)
        return self.__class__(other.data / self.data)

    def __rmatmul__(self, other: ArrayBase) -> ArrayBase:
        """Right-hand matrix multiplication."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if other.shape[1] != self.shape[0]:
            raise ValueError(f"{self.__class__.__name__} shapes are not aligned for matrix multiplication")
        return self.__class__(other.data @ self.data)

    # ---- in-place arithmetic operators ----

    def __iadd__(self, other: ArrayBase) -> ArrayBase:
        """In-place element-wise addition."""
        return self._inplace_elementwise(other, operator.iadd)

    def __isub__(self, other: ArrayBase) -> ArrayBase:
        """In-place element-wise subtraction."""
        return self._inplace_elementwise(other, operator.isub)

    def __imul__(self, other: ArrayBase) -> ArrayBase:
        """In-place element-wise multiplication."""
        return self._inplace_elementwise(other, operator.imul)

    def __itruediv__(self, other: ArrayBase) -> ArrayBase:
        """In-place element-wise division."""
        if not self._require_same_type(other):
            return NotImplemented
        self._require_same_shape(other)
        self._require_nonzero(other.data)
        self.data /= other.data
        return self

    def __imatmul__(self, other: ArrayBase) -> ArrayBase:
        """In-place matrix multiplication."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape[1] != other.shape[0]:
            raise ValueError(f"{self.__class__.__name__} shapes are not aligned for matrix multiplication")
        self.data = self.data @ other.data
        return self
