"""Ellipsoid class — the central data structure of the Ellipsoidal Toolbox.

An ellipsoid E(q, Q) in R^n is the set

    E = { x in R^n : (x - q)^T Q^{-1} (x - q) <= 1 }

where *q* is the center vector and *Q* is a symmetric positive
semi-definite shape (configuration) matrix.

Reference
---------
MATLAB Ellipsoidal Toolbox v3, ``@ellipsoid/ellipsoid.m``
Kurzhanskiy & Varaiya, "Ellipsoidal Toolbox", Technical Report
EECS-2006-46, UC Berkeley, 2006.
"""

from __future__ import annotations

import math
from typing import Optional

import numpy as np
from numpy.typing import ArrayLike

from .config import ABS_TOL
from .helpers import ell_inv, ell_simdiag, regularize

__all__ = ["Ellipsoid"]


class Ellipsoid:
    """An ellipsoid defined by a center vector and a shape matrix.

    Parameters
    ----------
    q : array_like or None
        Center vector of shape ``(n,)``.  If *None* and *Q* is given the
        center defaults to the origin.  If both are *None* an empty
        ellipsoid is created.
    Q : array_like or None
        Symmetric positive semi-definite shape matrix of shape ``(n, n)``.

    Raises
    ------
    ValueError
        If *Q* is not square, not symmetric, not positive semi-definite,
        or if the dimensions of *q* and *Q* do not match.

    Examples
    --------
    >>> E = Ellipsoid([1, 2], [[4, 0], [0, 1]])
    >>> E.center
    array([1., 2.])
    >>> E.rho(np.array([1.0, 0.0]))
    (3.0, array([3., 2.]))
    """

    # Give Ellipsoid higher priority than numpy arrays so that
    # A @ E and b + E dispatch to our __rmatmul__ / __radd__.
    __array_priority__ = 20.0

    # ------------------------------------------------------------------
    # Construction
    # ------------------------------------------------------------------

    def __init__(
        self,
        q: Optional[ArrayLike] = None,
        Q: Optional[ArrayLike] = None,
    ) -> None:
        if q is None and Q is None:
            self._center: Optional[np.ndarray] = None
            self._shape: Optional[np.ndarray] = None
            return

        if q is not None and Q is None:
            Q = q
            q = None

        Q_arr = np.asarray(Q, dtype=float)
        if Q_arr.ndim != 2 or Q_arr.shape[0] != Q_arr.shape[1]:
            raise ValueError("Shape matrix Q must be square.")

        n = Q_arr.shape[0]

        if not np.allclose(Q_arr, Q_arr.T, atol=ABS_TOL):
            raise ValueError("Shape matrix Q must be symmetric.")

        Q_arr = 0.5 * (Q_arr + Q_arr.T)

        min_eig = float(np.min(np.linalg.eigvalsh(Q_arr)))
        if min_eig < 0 and abs(min_eig) > ABS_TOL:
            raise ValueError(
                f"Shape matrix Q must be positive semi-definite "
                f"(smallest eigenvalue = {min_eig:.2e})."
            )

        if q is None:
            q_arr = np.zeros(n, dtype=float)
        else:
            q_arr = np.asarray(q, dtype=float).ravel()

        if q_arr.shape[0] != n:
            raise ValueError(
                f"Dimension mismatch: center has {q_arr.shape[0]} elements "
                f"but shape matrix is {n}x{n}."
            )

        self._center = q_arr
        self._shape = Q_arr

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def center(self) -> Optional[np.ndarray]:
        """Center vector *q*, or *None* for an empty ellipsoid."""
        return self._center

    @property
    def shape(self) -> Optional[np.ndarray]:
        """Shape matrix *Q*, or *None* for an empty ellipsoid."""
        return self._shape

    # ------------------------------------------------------------------
    # Basic queries
    # ------------------------------------------------------------------

    def parameters(self) -> tuple[Optional[np.ndarray], Optional[np.ndarray]]:
        """Return ``(q, Q)`` — the center and shape matrix.

        Returns
        -------
        q : ndarray or None
            Center vector.
        Q : ndarray or None
            Shape matrix.
        """
        return self._center, self._shape

    def dimension(self) -> tuple[int, int]:
        """Return the space dimension and the intrinsic dimension (rank).

        Returns
        -------
        space_dim : int
            Dimension of the ambient space (*n*).
        ell_dim : int
            Rank of the shape matrix.  Equals *space_dim* for
            non-degenerate ellipsoids.
        """
        if self.is_empty():
            return 0, 0
        Q = self._shape
        assert Q is not None
        n = Q.shape[0]
        return n, int(np.linalg.matrix_rank(Q, tol=ABS_TOL))

    def is_empty(self) -> bool:
        """Return *True* if this is an empty (uninitialised) ellipsoid."""
        return self._center is None or self._shape is None

    def is_degenerate(self) -> bool:
        """Return *True* if the shape matrix is rank-deficient.

        A degenerate ellipsoid is "flat" — it lives in a lower-dimensional
        subspace (e.g. an ellipse in 3-D space).
        """
        if self.is_empty():
            return False
        sd, ed = self.dimension()
        return ed < sd

    def max_eig(self) -> float:
        """Return the largest eigenvalue of the shape matrix.

        Raises
        ------
        ValueError
            If the ellipsoid is empty.
        """
        if self.is_empty():
            raise ValueError("Ellipsoid is empty.")
        return float(np.max(np.linalg.eigvalsh(self._shape)))

    def min_eig(self) -> float:
        """Return the smallest eigenvalue of the shape matrix.

        Raises
        ------
        ValueError
            If the ellipsoid is empty.
        """
        if self.is_empty():
            raise ValueError("Ellipsoid is empty.")
        return float(np.min(np.linalg.eigvalsh(self._shape)))

    def trace(self) -> float:
        """Return the trace of the shape matrix.

        Raises
        ------
        ValueError
            If the ellipsoid is empty.
        """
        if self.is_empty():
            raise ValueError("Ellipsoid is empty.")
        return float(np.trace(self._shape))

    def volume(self) -> float:
        """Return the n-dimensional volume of the ellipsoid.

        .. math::

            V = \\frac{\\pi^{n/2}}{\\Gamma(n/2 + 1)} \\sqrt{\\det Q}

        Returns 0 for degenerate ellipsoids.

        Raises
        ------
        ValueError
            If the ellipsoid is empty.
        """
        if self.is_empty():
            raise ValueError("Ellipsoid is empty.")
        assert self._shape is not None
        n = self._shape.shape[0]
        det_Q = np.linalg.det(self._shape)
        if det_Q < 0:
            return 0.0
        return (math.pi ** (n / 2) / math.gamma(n / 2 + 1)) * math.sqrt(det_Q)

    # ------------------------------------------------------------------
    # Support function  (THE core method)
    # ------------------------------------------------------------------

    def rho(
        self,
        L: ArrayLike,
    ) -> tuple[np.ndarray | float, np.ndarray]:
        """Evaluate the support function in one or more directions.

        The support function of E(q, Q) in direction *l* is

        .. math::

            \\rho(l \\mid \\mathcal{E}) = \\langle q, l \\rangle
            + \\sqrt{\\langle l, Q l \\rangle}

        and the boundary point where the maximum is attained is

        .. math::

            x^* = q + \\frac{Q l}{\\sqrt{\\langle l, Q l \\rangle}}

        Parameters
        ----------
        L : array_like
            Either a single direction vector of shape ``(n,)`` or a matrix
            whose **columns** are direction vectors, shape ``(n, k)``.

        Returns
        -------
        values : float or ndarray of shape ``(k,)``
            Support-function values.
        boundary_points : ndarray of shape ``(n,)`` or ``(n, k)``
            Boundary points of the ellipsoid corresponding to each
            direction.

        Raises
        ------
        ValueError
            If the ellipsoid is empty or dimensions do not match.
        """
        if self.is_empty():
            raise ValueError("Cannot compute support function of an empty ellipsoid.")

        q = self._center
        Q = self._shape
        assert q is not None and Q is not None
        n = q.shape[0]

        L_arr = np.asarray(L, dtype=float)
        single = L_arr.ndim == 1
        if single:
            L_arr = L_arr.reshape(-1, 1)

        if L_arr.shape[0] != n:
            raise ValueError(
                f"Direction dimension ({L_arr.shape[0]}) does not match "
                f"ellipsoid dimension ({n})."
            )

        k = L_arr.shape[1]
        Ql = Q @ L_arr                          # (n, k)
        lQl = np.sum(L_arr * Ql, axis=0)        # (k,)
        sr = np.sqrt(np.maximum(lQl, 0.0))       # (k,)

        sr_safe = np.where(sr == 0.0, np.finfo(float).eps, sr)

        values = q @ L_arr + sr                  # (k,)  broadcast q·l_i + sr_i
        boundary = q[:, np.newaxis] + Ql / sr_safe[np.newaxis, :]  # (n, k)

        if single:
            return float(values[0]), boundary[:, 0]
        return values, boundary

    # ------------------------------------------------------------------
    # Affine transformations
    # ------------------------------------------------------------------

    def linear_map(self, A: ArrayLike) -> Ellipsoid:
        """Apply a linear map: ``A @ E(q, Q) = E(Aq, A Q A^T)``.

        Parameters
        ----------
        A : array_like (m, n)
            Transformation matrix.

        Returns
        -------
        Ellipsoid
            The mapped ellipsoid in R^m.
        """
        if self.is_empty():
            return Ellipsoid()
        A_arr = np.asarray(A, dtype=float)
        if A_arr.ndim == 0:
            A_arr = A_arr.reshape(1, 1)
        q, Q = self._center, self._shape
        assert q is not None and Q is not None
        new_q = A_arr @ q
        new_Q = A_arr @ Q @ A_arr.T
        new_Q = 0.5 * (new_Q + new_Q.T)
        return Ellipsoid(new_q, new_Q)

    def __rmatmul__(self, A: ArrayLike) -> Ellipsoid:
        """Enable ``A @ E`` syntax for linear maps."""
        return self.linear_map(A)

    def __add__(self, b: ArrayLike) -> Ellipsoid:
        """Translate: ``E(q, Q) + b = E(q + b, Q)``.

        Parameters
        ----------
        b : array_like (n,)
            Translation vector.
        """
        if self.is_empty():
            return Ellipsoid()
        b_arr = np.asarray(b, dtype=float).ravel()
        assert self._center is not None and self._shape is not None
        return Ellipsoid(self._center + b_arr, self._shape.copy())

    def __radd__(self, b: ArrayLike) -> Ellipsoid:
        """Enable ``b + E`` syntax."""
        return self.__add__(b)

    def __sub__(self, b: ArrayLike) -> Ellipsoid:
        """Translate: ``E(q, Q) - b = E(q - b, Q)``.

        Parameters
        ----------
        b : array_like (n,)
            Translation vector.
        """
        if self.is_empty():
            return Ellipsoid()
        b_arr = np.asarray(b, dtype=float).ravel()
        assert self._center is not None and self._shape is not None
        return Ellipsoid(self._center - b_arr, self._shape.copy())

    def __neg__(self) -> Ellipsoid:
        """Negate center: ``-E(q, Q) = E(-q, Q)``."""
        if self.is_empty():
            return Ellipsoid()
        assert self._center is not None and self._shape is not None
        return Ellipsoid(-self._center, self._shape.copy())

    # ------------------------------------------------------------------
    # Geometry methods
    # ------------------------------------------------------------------

    def shape_transform(self, A: ArrayLike) -> Ellipsoid:
        """Modify only the shape matrix: ``E(q, Q) -> E(q, A Q A^T)``.

        Unlike :meth:`linear_map`, the center is preserved.

        Parameters
        ----------
        A : array_like (n, n) or scalar
            Transformation matrix or scalar multiplier.
        """
        if self.is_empty():
            return Ellipsoid()
        A_arr = np.asarray(A, dtype=float)
        q, Q = self._center, self._shape
        assert q is not None and Q is not None
        if A_arr.ndim == 0:
            new_Q = float(A_arr) ** 2 * Q
        else:
            new_Q = A_arr @ Q @ A_arr.T
        new_Q = 0.5 * (new_Q + new_Q.T)
        return Ellipsoid(q.copy(), new_Q)

    def inv(self) -> Ellipsoid:
        """Invert the shape matrix: ``E(q, Q) -> E(q, Q^{-1})``.

        Regularizes Q first if it is singular.
        """
        if self.is_empty():
            return Ellipsoid()
        q, Q = self._center, self._shape
        assert q is not None and Q is not None
        Q_inv = ell_inv(Q)
        Q_inv = 0.5 * (Q_inv + Q_inv.T)
        return Ellipsoid(q.copy(), Q_inv)

    def move_to_origin(self) -> Ellipsoid:
        """Return ellipsoid centered at the origin: ``E(q, Q) -> E(0, Q)``."""
        if self.is_empty():
            return Ellipsoid()
        assert self._shape is not None
        n = self._shape.shape[0]
        return Ellipsoid(np.zeros(n), self._shape.copy())

    def polar(self) -> Ellipsoid:
        """Compute the polar ellipsoid.

        The polar set of E(q, Q) is

        .. math::

            \\mathcal{E}^* = \\{ l \\in \\mathbb{R}^n \\mid
            \\rho(l \\mid \\mathcal{E}) \\le 1 \\}

        If the origin is an interior point of E(q, Q), its polar set
        is the ellipsoid E(w, W) where::

            M = (Q - q q^T)^{-1}
            w = -M q
            W = (1 + q^T M q) M

        Returns an empty ellipsoid if the origin is not interior or
        the ellipsoid is degenerate.
        """
        if self.is_empty() or self.is_degenerate():
            return Ellipsoid()
        q, Q = self._center, self._shape
        assert q is not None and Q is not None

        Q_inv = np.linalg.inv(Q)
        if q @ Q_inv @ q >= 1.0:
            return Ellipsoid()

        M = np.linalg.inv(Q - np.outer(q, q))
        M = 0.5 * (M + M.T)
        w = -M @ q
        W = (1.0 + q @ M @ q) * M
        W = 0.5 * (W + W.T)
        return Ellipsoid(w, W)

    def projection(self, B: ArrayLike) -> Ellipsoid:
        """Project the ellipsoid onto a subspace.

        .. math::

            \\text{proj}_B(\\mathcal{E}(q, Q)) =
            \\mathcal{E}(B^T q,\\; B^T Q B)

        Parameters
        ----------
        B : array_like (n, k)
            Matrix whose columns are orthogonal basis vectors of the
            target subspace.  Columns are normalised internally.
        """
        if self.is_empty():
            return Ellipsoid()
        q, Q = self._center, self._shape
        assert q is not None and Q is not None

        B_arr = np.asarray(B, dtype=float)
        n, k = B_arr.shape

        if n != Q.shape[0]:
            raise ValueError(
                f"Basis dimension ({n}) does not match "
                f"ellipsoid dimension ({Q.shape[0]})."
            )

        for i in range(k):
            nrm = np.linalg.norm(B_arr[:, i])
            if nrm < ABS_TOL:
                raise ValueError(f"Basis column {i} has near-zero norm.")
            B_arr[:, i] /= nrm

        for i in range(k):
            for j in range(i + 1, k):
                if abs(B_arr[:, i] @ B_arr[:, j]) > ABS_TOL:
                    raise ValueError("Basis columns must be orthogonal.")

        new_q = B_arr.T @ q
        new_Q = B_arr.T @ Q @ B_arr
        new_Q = 0.5 * (new_Q + new_Q.T)
        return Ellipsoid(new_q, new_Q)

    # ------------------------------------------------------------------
    # Point membership
    # ------------------------------------------------------------------

    def is_internal(
        self,
        X: ArrayLike,
    ) -> np.ndarray:
        """Check if points lie inside (or on the boundary of) this ellipsoid.

        A point *x* is internal to E(q, Q) when

        .. math::

            (x - q)^T Q^{-1} (x - q) \\le 1

        Degenerate ellipsoids are regularised before the check.

        Parameters
        ----------
        X : array_like
            Single point of shape ``(n,)`` or matrix whose **columns**
            are points, shape ``(n, k)``.

        Returns
        -------
        ndarray of bool, shape ``(k,)``
            ``True`` at positions where the point is internal.

        Raises
        ------
        ValueError
            If the ellipsoid is empty or dimensions do not match.
        """
        if self.is_empty():
            raise ValueError("Cannot check membership of an empty ellipsoid.")

        q, Q = self._center, self._shape
        assert q is not None and Q is not None
        n = q.shape[0]

        X_arr = np.asarray(X, dtype=float)
        single = X_arr.ndim == 1
        if single:
            X_arr = X_arr.reshape(-1, 1)

        if X_arr.shape[0] != n:
            raise ValueError(
                f"Point dimension ({X_arr.shape[0]}) does not match "
                f"ellipsoid dimension ({n})."
            )

        sd, rk = self.dimension()
        if rk < sd:
            Q = regularize(Q)

        Qi = ell_inv(Q)
        Qi = 0.5 * (Qi + Qi.T)

        diff = X_arr - q[:, np.newaxis]          # (n, k)
        vals = np.sum(diff * (Qi @ diff), axis=0)  # (k,)

        result = vals <= 1.0 + ABS_TOL
        if single:
            return result[0:1]
        return result

    # ------------------------------------------------------------------
    # Comparison operators
    # ------------------------------------------------------------------

    def is_bigger(self, other: Ellipsoid) -> bool:
        """Check if this ellipsoid would contain *other* with same center.

        Uses simultaneous diagonalization: if T diagonalises both
        shape matrices such that T Q1 T^T = I, then Q1 >= Q2 iff
        all diagonal entries of T Q2 T^T are <= 1.

        Parameters
        ----------
        other : Ellipsoid

        Returns
        -------
        bool
        """
        if self.is_empty() or other.is_empty():
            return False

        sd1, r1 = self.dimension()
        sd2, r2 = other.dimension()
        if sd1 != sd2:
            raise ValueError("Ellipsoids must have the same dimension.")
        if r1 < r2:
            return False

        A = self._shape
        B = other._shape
        assert A is not None and B is not None

        if r1 < sd1:
            A = regularize(A)

        T = ell_simdiag(A, B)
        diag_vals = np.diag(T @ B @ T.T)
        return bool(np.max(np.abs(diag_vals)) < 1.0 + ABS_TOL)

    def __gt__(self, other: Ellipsoid) -> bool:
        """``E1 > E2``: would E1 contain E2 if both centered at origin?"""
        if not isinstance(other, Ellipsoid):
            return NotImplemented
        return self.is_bigger(other)

    def __ge__(self, other: Ellipsoid) -> bool:
        """``E1 >= E2``: same as ``>``."""
        if not isinstance(other, Ellipsoid):
            return NotImplemented
        return self.is_bigger(other)

    def __lt__(self, other: Ellipsoid) -> bool:
        """``E1 < E2``: would E2 contain E1 if both centered at origin?"""
        if not isinstance(other, Ellipsoid):
            return NotImplemented
        return other.is_bigger(self)

    def __le__(self, other: Ellipsoid) -> bool:
        """``E1 <= E2``: same as ``<``."""
        if not isinstance(other, Ellipsoid):
            return NotImplemented
        return other.is_bigger(self)

    # ------------------------------------------------------------------
    # Dunder / display
    # ------------------------------------------------------------------

    def __repr__(self) -> str:
        if self.is_empty():
            return "Ellipsoid(empty)"
        n, r = self.dimension()
        return f"Ellipsoid(dim={n}, rank={r})"

    def __str__(self) -> str:
        if self.is_empty():
            return "Empty ellipsoid."
        n, r = self.dimension()
        lines = [
            f"Ellipsoid in R^{n} (rank {r}):",
            f"  center q = {self._center}",
            f"  shape  Q =\n{self._shape}",
        ]
        return "\n".join(lines)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Ellipsoid):
            return NotImplemented
        if self.is_empty() and other.is_empty():
            return True
        if self.is_empty() or other.is_empty():
            return False
        assert self._center is not None and other._center is not None
        assert self._shape is not None and other._shape is not None
        return bool(
            np.allclose(self._center, other._center, atol=ABS_TOL)
            and np.allclose(self._shape, other._shape, atol=ABS_TOL)
        )

    def __ne__(self, other: object) -> bool:
        result = self.__eq__(other)
        if result is NotImplemented:
            return NotImplemented
        return not result
