from __future__ import annotations

import numpy as np

from .array_base import ArrayBase
from .vector import Vector


class Matrix(ArrayBase):
    def __init__(self, data: np.ndarray) -> None:
        super().__init__(data)

    def _require_square(self, operation: str) -> None:
        if self.shape[0] != self.shape[1]:
            raise ValueError(f"{operation} is only defined for square matrices")

    # ---- matrix transforms ----

    def transpose(self) -> Matrix:
        """Return the transpose of the matrix."""
        return Matrix(self.data.T)

    @property
    def T(self) -> Matrix:
        """Return the transpose of the matrix."""
        return self.transpose()

    # ---- square-matrix operations ----

    def determinant(self) -> float:
        """Return the determinant of the matrix."""
        self._require_square("Determinant")
        return np.linalg.det(self.data)

    def inverse(self) -> Matrix:
        """Return the inverse of the matrix."""
        self._require_square("Inverse")
        return Matrix(np.linalg.inv(self.data))

    def trace(self) -> float:
        """Return the trace of the matrix."""
        self._require_square("Trace")
        return np.trace(self.data)

    def eigenvalues(self) -> Vector:
        """Return the eigenvalues of the matrix."""
        self._require_square("Eigenvalues")
        return Vector(np.linalg.eigvalsh(self.data))

    def eigenvectors(self) -> tuple[Vector, Matrix]:
        """Return the eigenvalues and eigenvectors of the matrix."""
        self._require_square("Eigenvectors")
        values, vectors = np.linalg.eigh(self.data)
        return Vector(values), Matrix(vectors)

    # ---- linear systems ----

    def solve(self, b: Vector) -> Vector:
        """Solve the linear system Ax = b for x."""
        self._require_square("Solving")
        return Vector(np.linalg.solve(self.data, b.data))