"""Unit tests for the Ellipsoid class.

Run with:
    python3 -m pytest tests/test_ellipsoid.py -v

Or simply:
    python3 tests/test_ellipsoid.py

Each test has a verbose docstring explaining WHAT is being tested and
WHY it matters mathematically.
"""

import math
import sys
import os

import numpy as np
import pytest

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from ellipsoidal_toolbox.core.ellipsoid import Ellipsoid
from ellipsoidal_toolbox.core.config import ABS_TOL


# ======================================================================
# 1. CONSTRUCTION & VALIDATION
# ======================================================================

class TestConstruction:
    """Test that ellipsoids are built correctly and bad inputs are rejected."""

    def test_basic_construction(self):
        """Create E([1,2], diag(4,1)) — a standard 2D ellipse.

        The center should be stored as a float64 array and the shape
        matrix should be symmetric.  This is the most basic "does it
        even work" test.
        """
        E = Ellipsoid([1, 2], [[4, 0], [0, 1]])
        assert E.center is not None
        np.testing.assert_array_almost_equal(E.center, [1.0, 2.0])
        np.testing.assert_array_almost_equal(E.shape, [[4, 0], [0, 1]])

    def test_shape_only_construction(self):
        """Passing only Q (no center) should default center to the origin.

        This matches the MATLAB behavior: E = ellipsoid(Q) centers at 0.
        """
        E = Ellipsoid(np.eye(3))
        np.testing.assert_array_almost_equal(E.center, [0.0, 0.0, 0.0])
        np.testing.assert_array_almost_equal(E.shape, np.eye(3))

    def test_empty_construction(self):
        """Ellipsoid() with no arguments creates an empty (uninitialised)
        ellipsoid.  This is the "null object" pattern — many operations
        return empty ellipsoids when results don't exist (e.g. empty
        intersection).
        """
        E = Ellipsoid()
        assert E.is_empty()
        assert E.center is None
        assert E.shape is None

    def test_1d_ellipsoid(self):
        """A 1D ellipsoid is just an interval.

        E(5, [[9]]) represents the interval [5 - 3, 5 + 3] = [2, 8].
        The shape matrix eigenvalue (9) is the squared half-length.
        """
        E = Ellipsoid([5], [[9]])
        assert E.dimension() == (1, 1)
        np.testing.assert_array_almost_equal(E.center, [5.0])

    def test_high_dimensional(self):
        """Ellipsoids in R^10 — verifying the toolbox scales to higher
        dimensions without issues.  This matters because the whole point
        of ellipsoidal calculus is that it scales well (O(n^2) storage)
        unlike polytopes.
        """
        n = 10
        q = np.zeros(n)
        Q = np.eye(n)
        E = Ellipsoid(q, Q)
        assert E.dimension() == (10, 10)

    def test_reject_nonsquare_shape(self):
        """Shape matrix must be square (n x n).  A 2x3 matrix makes no
        geometric sense — you can't define an ellipsoid with it.
        """
        with pytest.raises(ValueError, match="square"):
            Ellipsoid([0, 0], [[1, 0, 0], [0, 1, 0]])

    def test_reject_asymmetric_shape(self):
        """Shape matrix must be symmetric (Q = Q^T).  An asymmetric matrix
        would define different stretches depending on direction, which
        breaks the ellipsoid definition.
        """
        with pytest.raises(ValueError, match="symmetric"):
            Ellipsoid([0, 0], [[1, 5], [0, 1]])

    def test_reject_negative_definite(self):
        """Shape matrix must be positive semi-definite.  A negative
        eigenvalue would mean the "ellipsoid" extends in imaginary
        directions — geometrically meaningless.
        """
        with pytest.raises(ValueError, match="positive semi-definite"):
            Ellipsoid([0, 0], [[-1, 0], [0, 1]])

    def test_reject_dimension_mismatch(self):
        """Center must have the same number of elements as the shape
        matrix has rows/columns.  A 3D center with a 2x2 shape makes
        no sense.
        """
        with pytest.raises(ValueError, match="Dimension mismatch"):
            Ellipsoid([1, 2, 3], [[1, 0], [0, 1]])

    def test_nearly_symmetric_is_accepted(self):
        """Matrices that are symmetric up to floating-point noise should
        be accepted and symmetrised.  In real computations, you almost
        never get exactly Q == Q^T due to rounding.
        """
        Q = np.array([[4.0, 1.0 + 1e-15], [1.0, 4.0]])
        E = Ellipsoid([0, 0], Q)
        np.testing.assert_array_almost_equal(E.shape, E.shape.T)

    def test_nearly_zero_negative_eigenvalue_accepted(self):
        """Eigenvalues that are negative but within ABS_TOL of zero
        should be accepted.  This happens constantly in real
        computations — a matrix that should be PSD ends up with
        eigenvalue -1e-16 due to floating-point arithmetic.
        """
        Q = np.array([[1.0, 0.0], [0.0, -1e-14]])
        E = Ellipsoid([0, 0], Q)
        assert not E.is_empty()


# ======================================================================
# 2. PROPERTIES & BASIC QUERIES
# ======================================================================

class TestBasicQueries:
    """Test dimension, degeneracy, eigenvalue, trace, volume queries."""

    def test_dimension_full_rank(self):
        """A 3x3 identity shape matrix has rank 3 — full rank.

        space_dim=3 means we're in R^3.  ell_dim=3 means the ellipsoid
        fills all 3 dimensions (it's a solid ball, not a flat disk).
        """
        E = Ellipsoid(np.zeros(3), np.eye(3))
        sd, ed = E.dimension()
        assert sd == 3
        assert ed == 3

    def test_dimension_degenerate(self):
        """Shape diag(1, 1, 0) has rank 2 in R^3.

        This is a "flat" ellipsoid — a 2D disk floating in 3D space.
        space_dim=3 (we're in R^3) but ell_dim=2 (it only spans 2
        dimensions).  This comes up when projecting higher-dimensional
        ellipsoids or when a dynamical system has uncontrollable modes.
        """
        Q = np.diag([1.0, 1.0, 0.0])
        E = Ellipsoid(np.zeros(3), Q)
        sd, ed = E.dimension()
        assert sd == 3
        assert ed == 2

    def test_dimension_empty(self):
        """Empty ellipsoid has dimension (0, 0)."""
        E = Ellipsoid()
        assert E.dimension() == (0, 0)

    def test_is_degenerate_true(self):
        """rank(Q) < n means degenerate."""
        E = Ellipsoid([0, 0, 0], np.diag([1.0, 0.0, 1.0]))
        assert E.is_degenerate()

    def test_is_degenerate_false(self):
        """Full-rank shape is not degenerate."""
        E = Ellipsoid([0, 0], np.eye(2))
        assert not E.is_degenerate()

    def test_max_eig(self):
        """max_eig of diag(1, 4, 9) should be 9.

        This is the squared length of the longest semi-axis.  It tells
        you the "maximum extent" of the ellipsoid — important for
        bounding computations.
        """
        E = Ellipsoid(np.zeros(3), np.diag([1.0, 4.0, 9.0]))
        assert E.max_eig() == pytest.approx(9.0)

    def test_min_eig(self):
        """min_eig of diag(1, 4, 9) should be 1.

        The smallest eigenvalue controls whether the ellipsoid is
        "thin" in some direction — critical for degeneracy checks.
        """
        E = Ellipsoid(np.zeros(3), np.diag([1.0, 4.0, 9.0]))
        assert E.min_eig() == pytest.approx(1.0)

    def test_trace(self):
        """trace(diag(1, 4, 9)) = 14.

        The trace is the sum of squared semi-axis lengths.  It's a
        cheap measure of ellipsoid "size" (cheaper than volume).
        """
        E = Ellipsoid(np.zeros(3), np.diag([1.0, 4.0, 9.0]))
        assert E.trace() == pytest.approx(14.0)

    def test_volume_unit_circle(self):
        """Volume of the unit circle E(0, I_2) should be pi.

        In 2D, volume = pi * sqrt(det(Q)) = pi * sqrt(1) = pi.
        This is the basic sanity check for the volume formula.
        """
        E = Ellipsoid(np.zeros(2), np.eye(2))
        assert E.volume() == pytest.approx(math.pi, rel=1e-10)

    def test_volume_unit_sphere(self):
        """Volume of the unit sphere E(0, I_3) should be 4/3 * pi.

        In 3D, volume = (4/3)*pi*sqrt(det(I)) = 4*pi/3.
        """
        E = Ellipsoid(np.zeros(3), np.eye(3))
        expected = (4.0 / 3.0) * math.pi
        assert E.volume() == pytest.approx(expected, rel=1e-10)

    def test_volume_scaled_ellipse(self):
        """Volume of E(0, diag(4, 9)) in 2D should be pi * 2 * 3 = 6*pi.

        Semi-axes are sqrt(4)=2 and sqrt(9)=3.  Area = pi*a*b = 6*pi.
        General formula: pi * sqrt(det(Q)) = pi * sqrt(36) = 6*pi.
        """
        E = Ellipsoid(np.zeros(2), np.diag([4.0, 9.0]))
        expected = math.pi * 6.0
        assert E.volume() == pytest.approx(expected, rel=1e-10)

    def test_volume_degenerate_is_zero(self):
        """A degenerate ellipsoid has zero volume because it's flat.

        A 2D disk in 3D has zero 3D-volume (det(Q) = 0).
        """
        E = Ellipsoid(np.zeros(3), np.diag([1.0, 1.0, 0.0]))
        assert E.volume() == pytest.approx(0.0)

    def test_parameters(self):
        """parameters() should return the same (q, Q) that was passed in."""
        q_in = [1.0, 2.0]
        Q_in = [[4.0, 0.0], [0.0, 1.0]]
        E = Ellipsoid(q_in, Q_in)
        q, Q = E.parameters()
        np.testing.assert_array_almost_equal(q, q_in)
        np.testing.assert_array_almost_equal(Q, Q_in)

    def test_empty_queries_raise(self):
        """Calling max_eig/min_eig/trace/volume on an empty ellipsoid
        should raise ValueError — there's nothing to compute.
        """
        E = Ellipsoid()
        with pytest.raises(ValueError):
            E.max_eig()
        with pytest.raises(ValueError):
            E.min_eig()
        with pytest.raises(ValueError):
            E.trace()
        with pytest.raises(ValueError):
            E.volume()


# ======================================================================
# 3. SUPPORT FUNCTION  rho(l | E) = <q, l> + sqrt(<l, Q l>)
# ======================================================================

class TestSupportFunction:
    """Test the support function — the mathematical engine of the toolbox.

    The support function rho(l | E) answers: "how far does the
    ellipsoid E extend in direction l?"  It's defined as:

        rho(l | E) = <q, l> + sqrt(<l, Q*l>)

    where q is the center and Q is the shape matrix.

    The boundary point where this maximum is achieved is:

        x* = q + Q*l / sqrt(<l, Q*l>)
    """

    def test_axis_aligned_x(self):
        """rho([1,0] | E([1,2], diag(4,1))) = 1*1 + 0*2 + sqrt(1*4*1) = 3.

        Probing along the x-axis.  The ellipse stretches 2 units
        (sqrt(4)) from its center at x=1, so the furthest x-point is 3.
        The boundary point should be [3, 2] — at the rightmost tip.
        """
        E = Ellipsoid([1, 2], [[4, 0], [0, 1]])
        val, bp = E.rho(np.array([1.0, 0.0]))
        assert val == pytest.approx(3.0)
        np.testing.assert_array_almost_equal(bp, [3.0, 2.0])

    def test_axis_aligned_y(self):
        """rho([0,1] | E([1,2], diag(4,1))) = 0*1 + 1*2 + sqrt(0+1) = 3.

        Probing along the y-axis.  The ellipse stretches 1 unit from
        center at y=2, so furthest y-point is 3.
        Boundary point should be [1, 3] — at the top.
        """
        E = Ellipsoid([1, 2], [[4, 0], [0, 1]])
        val, bp = E.rho(np.array([0.0, 1.0]))
        assert val == pytest.approx(3.0)
        np.testing.assert_array_almost_equal(bp, [1.0, 3.0])

    def test_negative_direction(self):
        """rho([-1,0] | E([1,2], diag(4,1))) = -1 + 2 = 1.

        Probing in the -x direction.  The center contributes -1, and
        the shape extends sqrt(4)=2, so total is 1.
        Boundary point should be [-1, 2] — leftmost tip.
        """
        E = Ellipsoid([1, 2], [[4, 0], [0, 1]])
        val, bp = E.rho(np.array([-1.0, 0.0]))
        assert val == pytest.approx(1.0)
        np.testing.assert_array_almost_equal(bp, [-1.0, 2.0])

    def test_diagonal_direction(self):
        """rho([1,1] | E([0,0], I_2)) = 0 + sqrt(1+1) = sqrt(2).

        For the unit circle at origin, probing at 45 degrees.
        The boundary point should be [1/sqrt(2), 1/sqrt(2)] — the
        point on the unit circle at 45 degrees.
        """
        E = Ellipsoid(np.zeros(2), np.eye(2))
        val, bp = E.rho(np.array([1.0, 1.0]))
        assert val == pytest.approx(math.sqrt(2))
        s = 1.0 / math.sqrt(2)
        np.testing.assert_array_almost_equal(bp, [s, s])

    def test_unit_ball_support(self):
        """For the unit ball E(0, I_n), rho(l) = ||l|| for any direction.

        This is a fundamental property: the support function of a unit
        ball centered at origin is just the norm of the direction vector.
        """
        E = Ellipsoid(np.zeros(3), np.eye(3))
        l = np.array([3.0, 4.0, 0.0])
        val, bp = E.rho(l)
        assert val == pytest.approx(5.0)  # ||[3,4,0]|| = 5

    def test_center_shifts_support(self):
        """Translating the center by q shifts rho by <q, l>.

        If E1 = E(0, Q) and E2 = E(q, Q), then:
            rho(l | E2) = rho(l | E1) + <q, l>

        This is a basic linearity property of the support function.
        """
        Q = np.diag([4.0, 9.0])
        E_origin = Ellipsoid(np.zeros(2), Q)
        E_shifted = Ellipsoid([10.0, 20.0], Q)
        l = np.array([1.0, 0.0])

        val0, _ = E_origin.rho(l)
        val1, _ = E_shifted.rho(l)
        assert val1 == pytest.approx(val0 + 10.0)

    def test_multiple_directions(self):
        """Passing a matrix L of shape (n, k) should compute rho for
        all k directions at once and return arrays of shape (k,) and
        (n, k).

        This is the vectorised API matching the MATLAB rho(E, L) where
        L is a matrix whose columns are direction vectors.
        """
        E = Ellipsoid([1, 2], [[4, 0], [0, 1]])
        L = np.array([
            [1, 0, -1],
            [0, 1,  0],
        ], dtype=float)  # 3 directions as columns

        vals, bps = E.rho(L)
        assert vals.shape == (3,)
        assert bps.shape == (2, 3)

        assert vals[0] == pytest.approx(3.0)   # [1,0]
        assert vals[1] == pytest.approx(3.0)   # [0,1]
        assert vals[2] == pytest.approx(1.0)   # [-1,0]

    def test_support_of_scaled_ellipsoid(self):
        """E(0, alpha*I) has rho(l) = sqrt(alpha) * ||l||.

        Scaling the shape by alpha scales the support function by
        sqrt(alpha).  This is because semi-axes grow as sqrt(eigenvalues).
        """
        alpha = 16.0
        E = Ellipsoid(np.zeros(2), alpha * np.eye(2))
        l = np.array([3.0, 4.0])
        val, _ = E.rho(l)
        assert val == pytest.approx(math.sqrt(alpha) * 5.0)  # 4 * 5 = 20

    def test_3d_support(self):
        """3D test: E(0, diag(1, 4, 9)), rho([0, 0, 1]) = sqrt(9) = 3.

        The ellipsoid stretches 3 units along z.  Probing along z
        should give exactly 3.
        """
        E = Ellipsoid(np.zeros(3), np.diag([1.0, 4.0, 9.0]))
        val, bp = E.rho(np.array([0.0, 0.0, 1.0]))
        assert val == pytest.approx(3.0)
        np.testing.assert_array_almost_equal(bp, [0.0, 0.0, 3.0])

    def test_degenerate_direction(self):
        """When probing a degenerate ellipsoid in its "flat" direction,
        sqrt(<l, Ql>) = 0.

        E(0, diag(1, 0)) in direction [0, 1] gives rho = 0 + 0 = 0.
        The ellipsoid has zero extent in the y-direction.
        """
        E = Ellipsoid(np.zeros(2), np.diag([1.0, 0.0]))
        val, bp = E.rho(np.array([0.0, 1.0]))
        assert val == pytest.approx(0.0, abs=1e-10)

    def test_empty_raises(self):
        """Computing rho on an empty ellipsoid should raise ValueError."""
        E = Ellipsoid()
        with pytest.raises(ValueError, match="empty"):
            E.rho(np.array([1.0, 0.0]))

    def test_dimension_mismatch_raises(self):
        """Direction vector dimension must match ellipsoid dimension."""
        E = Ellipsoid(np.zeros(3), np.eye(3))
        with pytest.raises(ValueError, match="dimension"):
            E.rho(np.array([1.0, 0.0]))  # 2D direction for 3D ellipsoid

    def test_boundary_point_is_on_boundary(self):
        """The boundary point x* returned by rho must satisfy the
        ellipsoid equation:

            (x* - q)^T Q^{-1} (x* - q) = 1  (exactly on the boundary)

        This is the fundamental correctness check: if the point is
        NOT on the boundary, the support function implementation is
        wrong.
        """
        Q = np.array([[5.0, 1.0], [1.0, 3.0]])
        q = np.array([2.0, -1.0])
        E = Ellipsoid(q, Q)

        rng = np.random.default_rng(42)
        for _ in range(20):
            l = rng.standard_normal(2)
            _, bp = E.rho(l)
            diff = bp - q
            Q_inv = np.linalg.inv(Q)
            val = diff @ Q_inv @ diff
            assert val == pytest.approx(1.0, abs=1e-10), (
                f"Boundary point {bp} is not on the ellipsoid boundary "
                f"(got {val}, expected 1.0)"
            )


# ======================================================================
# 4. EQUALITY & DISPLAY
# ======================================================================

class TestEqualityAndDisplay:
    """Test __eq__, __ne__, __repr__, __str__."""

    def test_equal_ellipsoids(self):
        """Two ellipsoids with identical (q, Q) should be equal."""
        E1 = Ellipsoid([1, 2], [[4, 0], [0, 1]])
        E2 = Ellipsoid([1, 2], [[4, 0], [0, 1]])
        assert E1 == E2

    def test_different_centers(self):
        """Different centers => not equal."""
        E1 = Ellipsoid([1, 2], np.eye(2))
        E2 = Ellipsoid([3, 4], np.eye(2))
        assert E1 != E2

    def test_different_shapes(self):
        """Different shape matrices => not equal."""
        E1 = Ellipsoid(np.zeros(2), np.eye(2))
        E2 = Ellipsoid(np.zeros(2), 2 * np.eye(2))
        assert E1 != E2

    def test_both_empty_are_equal(self):
        """Two empty ellipsoids should be equal."""
        assert Ellipsoid() == Ellipsoid()

    def test_empty_vs_nonempty(self):
        """Empty != non-empty."""
        assert Ellipsoid() != Ellipsoid(np.zeros(2), np.eye(2))

    def test_nearly_equal(self):
        """Ellipsoids that differ by less than ABS_TOL should be equal.

        This tests numerical robustness — in real computations you
        never get exact floating-point equality.
        """
        E1 = Ellipsoid([1, 2], [[4, 0], [0, 1]])
        q2 = np.array([1.0 + 1e-10, 2.0 - 1e-10])
        E2 = Ellipsoid(q2, [[4, 0], [0, 1]])
        assert E1 == E2

    def test_repr_nonempty(self):
        """repr should show dim and rank."""
        E = Ellipsoid(np.zeros(3), np.eye(3))
        assert "dim=3" in repr(E)
        assert "rank=3" in repr(E)

    def test_repr_empty(self):
        """repr of empty should say 'empty'."""
        assert "empty" in repr(Ellipsoid())

    def test_str_contains_center(self):
        """str should show the center values."""
        E = Ellipsoid([1, 2], np.eye(2))
        s = str(E)
        assert "center" in s
        assert "shape" in s

    def test_ne_with_non_ellipsoid(self):
        """Comparing with a non-Ellipsoid should return NotImplemented
        (Python will then fall back to default object comparison).
        """
        E = Ellipsoid(np.zeros(2), np.eye(2))
        assert E != "not an ellipsoid"


# ======================================================================
# Main entry point (for running without pytest)
# ======================================================================

if __name__ == "__main__":
    pytest.main([__file__, "-v", "--tb=short"])
