from __future__ import annotations
import numpy as np
import copy

_RTOL = 1e-9
_ATOL = 1e-12

class ArrayBase:
    __slots__ = ['data']
    
    def __init__(self, data) -> 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):
        return self.data.shape

    @property
    def ndim(self):
        return self.data.ndim

    @property
    def size(self):
        return self.data.size

    @property
    def dtype(self):
        return self.data.dtype

    # ---- representation ----

    def __repr__(self):
        return f"{self.__class__.__name__}({self.data})"
    
    def __str__(self):
        return str(self.data)
    
    def __bool__(self):
        """Return True if the array is not empty."""
        return self.size > 0
    
    # ---- length and iteration ----
    
    def __len__(self):
        """Return the size of the first dimension."""
        return self.data.shape[0]

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

    def __contains__(self, item):
        """Check if item is in the array."""
        if np.isscalar(item):
            return bool(np.any(np.isclose(self.data, item, rtol=_RTOL, atol=_ATOL)))
        return item in self.data
    
    def __setitem__(self, key, value):
        """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 scalar if result is 0-dimensional, otherwise wrap in the same class
        if np.isscalar(result) or result.ndim == 0:
            return result
        return self.__class__(result)

    # ---- comparison operators ----

    def __eq__(self, other):
        """Element-wise equality comparison."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        return self.__class__(np.isclose(self.data, other.data, rtol=_RTOL, atol=_ATOL))

    def __ne__(self, other):
        """Element-wise inequality comparison."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        return self.__class__(np.logical_not(np.isclose(self.data, other.data, rtol=_RTOL, atol=_ATOL)))

    def __lt__(self, other):
        """Element-wise less than comparison."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        return self.__class__(self.data < other.data)

    def __le__(self, other):
        """Element-wise less than or equal comparison."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        return self.__class__(self.data <= other.data)

    def __gt__(self, other):
        """Element-wise greater than comparison."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        return self.__class__(self.data > other.data)

    def __ge__(self, other):
        """Element-wise greater than or equal comparison."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        return self.__class__(self.data >= other.data)

    # ---- unary operators ----

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

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

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

    # ---- copy methods ----

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

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

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

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

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

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

    def sum(self, axis=None, keepdims=False):
        """Sum of array elements over a given axis."""
        result = self.data.sum(axis=axis, keepdims=keepdims)
        if np.isscalar(result):
            return result
        return self.__class__(result)

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

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

    def astype(self, dtype):
        """Return array with specified dtype."""
        return self.__class__(self.data.astype(dtype))
    
    # arithmetic ops
    def __add__(self, other):
        """Element-wise addition."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape != other.shape:
            raise ValueError(f"{self.__class__.__name__} must have the same shape for addition")
        return self.__class__(self.data + other.data)
    
    def __sub__(self, other):
        """Element-wise subtraction."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape != other.shape:
            raise ValueError(f"{self.__class__.__name__} must have the same shape for subtraction")
        return self.__class__(self.data - other.data)
    
    def __mul__(self, other):
        """Element-wise multiplication."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape != other.shape:
            raise ValueError(f"{self.__class__.__name__} must have the same shape for multiplication")
        return self.__class__(self.data * other.data)
    
    def __truediv__(self, other):
        """Element-wise division."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape != other.shape:
            raise ValueError(f"{self.__class__.__name__} must have the same shape for division")
        if np.any(np.isclose(other.data, 0.0, rtol=_RTOL, atol=_ATOL)):
            raise ZeroDivisionError("Cannot divide by zero")
        return self.__class__(self.data / other.data)
    
    def __matmul__(self, other):
        """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")
        return self.__class__(self.data @ other.data)
    
    # right-hand versions of arithmetic ops
    def __radd__(self, other):
        """Right-hand element-wise addition."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape != other.shape:
            raise ValueError(f"{self.__class__.__name__} must have the same shape for addition")
        return self.__class__(other.data + self.data)
    
    def __rsub__(self, other):
        """Right-hand element-wise subtraction."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape != other.shape:
            raise ValueError(f"{self.__class__.__name__} must have the same shape for subtraction")
        return self.__class__(other.data - self.data)
    
    def __rmul__(self, other):
        """Right-hand element-wise multiplication."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape != other.shape:
            raise ValueError(f"{self.__class__.__name__} must have the same shape for multiplication")
        return self.__class__(other.data * self.data)
    
    def __rtruediv__(self, other):
        """Right-hand element-wise division."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape != other.shape:
            raise ValueError(f"{self.__class__.__name__} must have the same shape for division")
        if np.any(np.isclose(self.data, 0.0, rtol=_RTOL, atol=_ATOL)):
            raise ZeroDivisionError("Cannot divide by zero")
        return self.__class__(other.data / self.data)
    
    def __rmatmul__(self, other):
        """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 versions of arithmetic ops
    def __iadd__(self, other):
        """In-place element-wise addition."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape != other.shape:
            raise ValueError(f"{self.__class__.__name__} must have the same shape for addition")
        self.data += other.data
        return self
    
    def __isub__(self, other):
        """In-place element-wise subtraction."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape != other.shape:
            raise ValueError(f"{self.__class__.__name__} must have the same shape for subtraction")
        self.data -= other.data
        return self
    
    def __imul__(self, other):
        """In-place element-wise multiplication."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape != other.shape:
            raise ValueError(f"{self.__class__.__name__} must have the same shape for multiplication")
        self.data *= other.data
        return self
    
    def __itruediv__(self, other):
        """In-place element-wise division."""
        if not isinstance(other, self.__class__):
            return NotImplemented
        if self.shape != other.shape:
            raise ValueError(f"{self.__class__.__name__} must have the same shape for division")
        if np.any(np.isclose(other.data, 0.0, rtol=_RTOL, atol=_ATOL)):
            raise ZeroDivisionError("Cannot divide by zero")
        self.data /= other.data
        return self
    
    def __imatmul__(self, other):
        """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
    