"""Unit tests for Phase 2A: Operator overloads, geometry methods, comparisons.

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

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.helpers import regularize, ell_valign, ell_simdiag, ell_inv
from ellipsoidal_toolbox.core.config import ABS_TOL


# ======================================================================
# 1. HELPER FUNCTIONS
# ======================================================================

class TestHelpers:
    """Test the low-level helper functions in core/helpers.py."""

    def test_regularize(self):
        """regularize(Q) should add delta*I to a singular matrix,
        making it full-rank while barely changing it.
        """
        Q = np.diag([1.0, 0.0])
        R = regularize(Q)
        assert np.linalg.matrix_rank(R) == 2
        assert R[0, 0] == pytest.approx(1.0 + ABS_TOL)
        assert R[1, 1] == pytest.approx(ABS_TOL)

    def test_ell_inv_full_rank(self):
        """ell_inv of a full-rank matrix should be the standard inverse."""
        Q = np.array([[4.0, 1.0], [1.0, 3.0]])
        Q_inv = ell_inv(Q)
        np.testing.assert_array_almost_equal(Q @ Q_inv, np.eye(2), decimal=10)

    def test_ell_inv_singular_regularizes(self):
        """ell_inv of a singular matrix should regularize first, then invert.
        The result times the regularized Q should give ~identity.
        """
        Q = np.diag([1.0, 0.0])
        Q_inv = ell_inv(Q)
        assert Q_inv.shape == (2, 2)
        assert np.isfinite(Q_inv).all()

    def test_ell_valign_parallel_vectors(self):
        """ell_valign of two parallel vectors should return ~identity
        (or a rotation that maps one to the other's direction).
        """
        v1 = np.array([1.0, 0.0, 0.0])
        v2 = np.array([3.0, 0.0, 0.0])
        S = ell_valign(v1, v2)
        np.testing.assert_array_almost_equal(S @ v1 / np.linalg.norm(v1),
                                              v2 / np.linalg.norm(v2))

    def test_ell_valign_orthogonal(self):
        """S from ell_valign must be an orthogonal matrix (S^T S = I).
        This preserves lengths and angles.
        """
        rng = np.random.default_rng(42)
        for _ in range(10):
            v1 = rng.standard_normal(4)
            v2 = rng.standard_normal(4)
            S = ell_valign(v1, v2)
            np.testing.assert_array_almost_equal(S.T @ S, np.eye(4), decimal=10)

    def test_ell_valign_result_is_parallel(self):
        """S v1 should be parallel to v2 (their cross product ~ 0 in 3D,
        or more generally the angle between them ~ 0).
        """
        rng = np.random.default_rng(99)
        for _ in range(10):
            v1 = rng.standard_normal(3)
            v2 = rng.standard_normal(3)
            S = ell_valign(v1, v2)
            Sv1 = S @ v1
            cos_angle = (Sv1 @ v2) / (np.linalg.norm(Sv1) * np.linalg.norm(v2))
            assert abs(abs(cos_angle) - 1.0) < 1e-10

    def test_ell_simdiag(self):
        """After simultaneous diagonalization T of (A, B):
        T A T^T should be ~identity, and T B T^T should be diagonal.
        """
        A = np.array([[5.0, 1.0], [1.0, 3.0]])
        B = np.array([[2.0, 0.5], [0.5, 1.0]])
        T = ell_simdiag(A, B)
        TATT = T @ A @ T.T
        TBTT = T @ B @ T.T
        np.testing.assert_array_almost_equal(TATT, np.eye(2), decimal=8)
        off_diag = TBTT - np.diag(np.diag(TBTT))
        assert np.max(np.abs(off_diag)) < 1e-8


# ======================================================================
# 2. LINEAR MAP  (A @ E)
# ======================================================================

class TestLinearMap:
    """Test A @ E = E(Aq, AQA^T) — the fundamental operation for
    propagating ellipsoids through linear dynamics.
    """

    def test_identity_map(self):
        """I @ E should return the same ellipsoid.

        The identity transformation changes nothing.
        """
        E = Ellipsoid([1, 2], [[4, 0], [0, 1]])
        result = np.eye(2) @ E
        assert result == E

    def test_scaling(self):
        """(2*I) @ E(q, Q) = E(2q, 4Q).

        Scaling by 2 doubles the center and quadruples the shape
        (since AQA^T = 2*Q*2 = 4Q).
        """
        E = Ellipsoid([1, 0], np.eye(2))
        result = (2 * np.eye(2)) @ E
        np.testing.assert_array_almost_equal(result.center, [2.0, 0.0])
        np.testing.assert_array_almost_equal(result.shape, 4 * np.eye(2))

    def test_rotation(self):
        """Rotating a unit circle should give back a unit circle.

        A rotation preserves all eigenvalues (the shape doesn't change
        for a sphere), but the center rotates.
        """
        theta = np.pi / 4
        R = np.array([[np.cos(theta), -np.sin(theta)],
                       [np.sin(theta),  np.cos(theta)]])
        E = Ellipsoid([1, 0], np.eye(2))
        result = R @ E
        np.testing.assert_array_almost_equal(result.shape, np.eye(2), decimal=10)
        expected_center = R @ np.array([1.0, 0.0])
        np.testing.assert_array_almost_equal(result.center, expected_center)

    def test_projection_via_map(self):
        """A non-square matrix maps E from R^3 to R^2.

        This is how you extract the (x, y) components of a 3D ellipsoid.
        """
        E = Ellipsoid(np.zeros(3), np.eye(3))
        A = np.array([[1, 0, 0], [0, 1, 0]])  # project onto x-y
        result = A @ E
        sd, _ = result.dimension()
        assert sd == 2
        np.testing.assert_array_almost_equal(result.shape, np.eye(2))

    def test_volume_scales_correctly(self):
        """vol(A @ E) = |det(A)| * vol(E) for square A.

        The determinant tells you how volume scales under linear maps.
        """
        E = Ellipsoid(np.zeros(2), np.eye(2))
        A = np.array([[2, 0], [0, 3]])
        result = A @ E
        assert result.volume() == pytest.approx(6.0 * E.volume(), rel=1e-10)

    def test_empty_returns_empty(self):
        """Linear map of an empty ellipsoid should be empty."""
        result = np.eye(2) @ Ellipsoid()
        assert result.is_empty()

    def test_linear_map_method(self):
        """E.linear_map(A) should give the same result as A @ E."""
        E = Ellipsoid([1, 2], [[4, 0], [0, 1]])
        A = np.array([[1, 2], [0, 1]])
        assert (A @ E) == E.linear_map(A)


# ======================================================================
# 3. TRANSLATION  (E + b, E - b)
# ======================================================================

class TestTranslation:
    """Test E + b = E(q+b, Q) and E - b = E(q-b, Q).

    Translation moves the center without changing the shape.
    """

    def test_add_vector(self):
        """E([0,0], I) + [3,4] should have center [3,4] and same shape."""
        E = Ellipsoid(np.zeros(2), np.eye(2))
        result = E + np.array([3.0, 4.0])
        np.testing.assert_array_almost_equal(result.center, [3, 4])
        np.testing.assert_array_almost_equal(result.shape, np.eye(2))

    def test_sub_vector(self):
        """E([5,5], I) - [3,4] should have center [2,1]."""
        E = Ellipsoid([5, 5], np.eye(2))
        result = E - np.array([3.0, 4.0])
        np.testing.assert_array_almost_equal(result.center, [2, 1])

    def test_add_then_sub_is_identity(self):
        """(E + b) - b should give back the original ellipsoid."""
        E = Ellipsoid([1, 2], [[4, 0], [0, 1]])
        b = np.array([10.0, -5.0])
        result = (E + b) - b
        assert result == E

    def test_radd(self):
        """b + E should work the same as E + b."""
        E = Ellipsoid(np.zeros(2), np.eye(2))
        b = np.array([1.0, 2.0])
        r1 = E + b
        r2 = b + E
        assert r1 == r2

    def test_volume_unchanged(self):
        """Translation does not change volume."""
        E = Ellipsoid(np.zeros(2), np.diag([4.0, 9.0]))
        shifted = E + np.array([100.0, 200.0])
        assert shifted.volume() == pytest.approx(E.volume())

    def test_empty_add(self):
        """Adding to empty gives empty."""
        assert (Ellipsoid() + np.array([1.0])).is_empty()


# ======================================================================
# 4. NEGATION  (-E)
# ======================================================================

class TestNegation:
    """Test -E(q, Q) = E(-q, Q)."""

    def test_negate_center(self):
        """-E should flip the center and keep the shape."""
        E = Ellipsoid([3, -4], [[4, 0], [0, 1]])
        neg = -E
        np.testing.assert_array_almost_equal(neg.center, [-3, 4])
        np.testing.assert_array_almost_equal(neg.shape, E.shape)

    def test_double_negate(self):
        """-(-E) should give back the original."""
        E = Ellipsoid([1, 2], np.eye(2))
        assert -(-E) == E

    def test_negate_empty(self):
        """-empty should be empty."""
        assert (-Ellipsoid()).is_empty()


# ======================================================================
# 5. SHAPE TRANSFORM
# ======================================================================

class TestShapeTransform:
    """Test shape_transform(A): E(q, Q) -> E(q, AQA^T).

    Unlike linear_map (A @ E), the center is NOT transformed.
    """

    def test_center_preserved(self):
        """After shape_transform, center should remain unchanged."""
        E = Ellipsoid([5, 10], np.eye(2))
        result = E.shape_transform(2.0 * np.eye(2))
        np.testing.assert_array_almost_equal(result.center, [5, 10])
        np.testing.assert_array_almost_equal(result.shape, 4 * np.eye(2))

    def test_scalar_scaling(self):
        """shape_transform(scalar) should scale Q by scalar^2."""
        E = Ellipsoid(np.zeros(2), np.eye(2))
        result = E.shape_transform(3.0)
        np.testing.assert_array_almost_equal(result.shape, 9 * np.eye(2))

    def test_vs_linear_map(self):
        """shape_transform and linear_map should produce same shape but
        different centers (linear_map also transforms center).
        """
        A = np.array([[2, 0], [0, 3]])
        E = Ellipsoid([1, 1], np.eye(2))
        mapped = A @ E
        shaped = E.shape_transform(A)
        np.testing.assert_array_almost_equal(mapped.shape, shaped.shape)
        assert not np.allclose(mapped.center, shaped.center)

    def test_empty(self):
        """shape_transform on empty gives empty."""
        assert Ellipsoid().shape_transform(np.eye(2)).is_empty()


# ======================================================================
# 6. INV
# ======================================================================

class TestInv:
    """Test inv(): E(q, Q) -> E(q, Q^{-1})."""

    def test_identity_inv(self):
        """inv(E(q, I)) should give E(q, I) since I^{-1} = I."""
        E = Ellipsoid([1, 2], np.eye(2))
        result = E.inv()
        assert result == E

    def test_diagonal_inv(self):
        """inv(diag(4, 9)) = diag(1/4, 1/9)."""
        E = Ellipsoid(np.zeros(2), np.diag([4.0, 9.0]))
        result = E.inv()
        np.testing.assert_array_almost_equal(result.shape, np.diag([0.25, 1.0/9.0]))

    def test_center_preserved(self):
        """inv should not change the center."""
        E = Ellipsoid([5, 10], np.diag([4.0, 9.0]))
        result = E.inv()
        np.testing.assert_array_almost_equal(result.center, [5, 10])

    def test_double_inv(self):
        """inv(inv(E)) should give back the original shape."""
        Q = np.array([[5.0, 1.0], [1.0, 3.0]])
        E = Ellipsoid(np.zeros(2), Q)
        result = E.inv().inv()
        np.testing.assert_array_almost_equal(result.shape, Q, decimal=10)

    def test_empty(self):
        """inv of empty gives empty."""
        assert Ellipsoid().inv().is_empty()


# ======================================================================
# 7. MOVE TO ORIGIN
# ======================================================================

class TestMoveToOrigin:
    """Test move_to_origin: E(q, Q) -> E(0, Q)."""

    def test_basic(self):
        """Center should become zero, shape unchanged."""
        E = Ellipsoid([100, 200], np.diag([4.0, 9.0]))
        result = E.move_to_origin()
        np.testing.assert_array_almost_equal(result.center, [0, 0])
        np.testing.assert_array_almost_equal(result.shape, E.shape)

    def test_already_at_origin(self):
        """Moving origin-centered ellipsoid should give the same thing."""
        E = Ellipsoid(np.zeros(3), np.eye(3))
        assert E.move_to_origin() == E

    def test_empty(self):
        """move_to_origin on empty gives empty."""
        assert Ellipsoid().move_to_origin().is_empty()


# ======================================================================
# 8. POLAR ELLIPSOID
# ======================================================================

class TestPolar:
    """Test the polar ellipsoid.

    The polar of E(q, Q) is { l : rho(l | E) <= 1 }.
    For origin-centered unit ball E(0, I), the polar is also E(0, I).
    """

    def test_unit_ball_polar(self):
        """The polar of the unit ball at origin is the unit ball itself.

        rho(l | E(0, I)) = ||l||, so { l : ||l|| <= 1 } is E(0, I).
        """
        E = Ellipsoid(np.zeros(3), np.eye(3))
        P = E.polar()
        np.testing.assert_array_almost_equal(P.center, [0, 0, 0])
        np.testing.assert_array_almost_equal(P.shape, np.eye(3), decimal=10)

    def test_scaled_ball_polar(self):
        """For E(0, alpha*I), polar should be E(0, (1/alpha)*I).

        rho(l | E(0, aI)) = sqrt(a)*||l||, so { l : sqrt(a)*||l|| <= 1 }
        is { l : ||l|| <= 1/sqrt(a) } = E(0, (1/a)*I).
        """
        alpha = 4.0
        E = Ellipsoid(np.zeros(2), alpha * np.eye(2))
        P = E.polar()
        np.testing.assert_array_almost_equal(P.center, [0, 0])
        np.testing.assert_array_almost_equal(P.shape, (1.0/alpha) * np.eye(2), decimal=10)

    def test_origin_not_interior_returns_empty(self):
        """If origin is outside the ellipsoid, polar is empty.

        E([10, 0], I) has center far from origin. Origin is not interior.
        """
        E = Ellipsoid([10, 0], np.eye(2))
        P = E.polar()
        assert P.is_empty()

    def test_degenerate_returns_empty(self):
        """Degenerate ellipsoid has no well-defined polar."""
        E = Ellipsoid(np.zeros(3), np.diag([1.0, 1.0, 0.0]))
        P = E.polar()
        assert P.is_empty()

    def test_empty_returns_empty(self):
        """polar of empty is empty."""
        assert Ellipsoid().polar().is_empty()

    def test_polar_with_offset_center(self):
        """E(q, Q) with origin just barely inside — verify the
        polar formulas produce a valid (PSD) shape matrix.
        """
        E = Ellipsoid([0.1, 0.0], 2.0 * np.eye(2))
        P = E.polar()
        assert not P.is_empty()
        assert P.min_eig() >= -ABS_TOL


# ======================================================================
# 9. PROJECTION
# ======================================================================

class TestProjection:
    """Test projection onto a subspace."""

    def test_project_3d_to_2d(self):
        """Project E(0, I_3) onto x-y plane.

        Result should be E(0, I_2) — the unit disk.
        """
        E = Ellipsoid(np.zeros(3), np.eye(3))
        B = np.array([[1, 0], [0, 1], [0, 0]], dtype=float)
        P = E.projection(B)
        sd, _ = P.dimension()
        assert sd == 2
        np.testing.assert_array_almost_equal(P.center, [0, 0])
        np.testing.assert_array_almost_equal(P.shape, np.eye(2))

    def test_project_3d_to_1d(self):
        """Project E(0, diag(4, 9, 16)) onto z-axis.

        Result should be E(0, [[16]]) — an interval of half-length 4.
        """
        E = Ellipsoid(np.zeros(3), np.diag([4.0, 9.0, 16.0]))
        B = np.array([[0], [0], [1]], dtype=float)
        P = E.projection(B)
        assert P.dimension() == (1, 1)
        np.testing.assert_array_almost_equal(P.shape, [[16.0]])

    def test_projection_preserves_center_component(self):
        """Center should project correctly: B^T q."""
        E = Ellipsoid([1, 2, 3], np.eye(3))
        B = np.array([[1, 0], [0, 1], [0, 0]], dtype=float)
        P = E.projection(B)
        np.testing.assert_array_almost_equal(P.center, [1, 2])

    def test_non_orthogonal_basis_raises(self):
        """Basis columns must be orthogonal."""
        E = Ellipsoid(np.zeros(3), np.eye(3))
        B = np.array([[1, 1], [0, 1], [0, 0]], dtype=float)
        with pytest.raises(ValueError, match="orthogonal"):
            E.projection(B)

    def test_empty(self):
        """Projecting empty gives empty."""
        B = np.array([[1, 0], [0, 1], [0, 0]], dtype=float)
        assert Ellipsoid().projection(B).is_empty()


# ======================================================================
# 10. IS_BIGGER / COMPARISON OPERATORS
# ======================================================================

class TestIsBigger:
    """Test is_bigger and the comparison operators (>, >=, <, <=).

    is_bigger(E1, E2) checks: would E1 contain E2 if both were
    centered at the origin?  This is determined by comparing the
    shape matrices via simultaneous diagonalization.
    """

    def test_larger_contains_smaller(self):
        """E(0, 4I) should be bigger than E(0, I).

        The first has semi-axes of length 2, the second has length 1.
        """
        E1 = Ellipsoid(np.zeros(2), 4 * np.eye(2))
        E2 = Ellipsoid(np.zeros(2), np.eye(2))
        assert E1.is_bigger(E2)
        assert not E2.is_bigger(E1)

    def test_equal_is_bigger(self):
        """An ellipsoid should be bigger-or-equal to itself."""
        E = Ellipsoid(np.zeros(2), np.eye(2))
        assert E.is_bigger(E)

    def test_one_axis_smaller(self):
        """diag(4, 1) vs diag(1, 4): neither contains the other.

        First is wide (x), second is tall (y).
        """
        E1 = Ellipsoid(np.zeros(2), np.diag([4.0, 1.0]))
        E2 = Ellipsoid(np.zeros(2), np.diag([1.0, 4.0]))
        assert not E1.is_bigger(E2)
        assert not E2.is_bigger(E1)

    def test_gt_operator(self):
        """The > operator should delegate to is_bigger."""
        E1 = Ellipsoid(np.zeros(2), 4 * np.eye(2))
        E2 = Ellipsoid(np.zeros(2), np.eye(2))
        assert E1 > E2
        assert not E2 > E1

    def test_lt_operator(self):
        """E2 < E1 should be the same as E1 > E2."""
        E1 = Ellipsoid(np.zeros(2), 4 * np.eye(2))
        E2 = Ellipsoid(np.zeros(2), np.eye(2))
        assert E2 < E1
        assert not E1 < E2

    def test_ge_le_operators(self):
        """>= and <= should work like > and <."""
        E1 = Ellipsoid(np.zeros(2), 4 * np.eye(2))
        E2 = Ellipsoid(np.zeros(2), np.eye(2))
        assert E1 >= E2
        assert E2 <= E1

    def test_centers_ignored(self):
        """is_bigger only compares shapes, not centers.

        Even if centers are different, the comparison should work
        (it's "would contain if same center").
        """
        E1 = Ellipsoid([100, 200], 4 * np.eye(2))
        E2 = Ellipsoid([-50, -50], np.eye(2))
        assert E1.is_bigger(E2)

    def test_dimension_mismatch_raises(self):
        """Comparing ellipsoids of different dimensions should raise."""
        E1 = Ellipsoid(np.zeros(2), np.eye(2))
        E2 = Ellipsoid(np.zeros(3), np.eye(3))
        with pytest.raises(ValueError, match="same dimension"):
            E1.is_bigger(E2)

    def test_3d_comparison(self):
        """3D test: diag(9, 4, 1) should contain diag(4, 1, 1)."""
        E1 = Ellipsoid(np.zeros(3), np.diag([9.0, 4.0, 1.0]))
        E2 = Ellipsoid(np.zeros(3), np.diag([4.0, 1.0, 1.0]))
        assert E1.is_bigger(E2)

    def test_non_diagonal(self):
        """Test with off-diagonal shape matrices (rotated ellipsoids).

        Q1 = [[5,1],[1,3]] has eigenvalues ~5.41 and ~2.59.
        Q2 = I has eigenvalues 1 and 1.
        So Q1 should contain Q2.
        """
        Q1 = np.array([[5.0, 1.0], [1.0, 3.0]])
        E1 = Ellipsoid(np.zeros(2), Q1)
        E2 = Ellipsoid(np.zeros(2), np.eye(2))
        assert E1.is_bigger(E2)


# ======================================================================
# 11. INTEGRATION TESTS
# ======================================================================

class TestIntegration:
    """Combined operations testing multiple Phase 2A features together."""

    def test_rotate_translate_support(self):
        """Rotate, then translate, then check support function.

        For E = R @ E0 + b, we have rho(l | E) = <l, b> + rho(R^T l | E0).
        """
        theta = np.pi / 3
        R = np.array([[np.cos(theta), -np.sin(theta)],
                       [np.sin(theta),  np.cos(theta)]])
        b = np.array([5.0, 0.0])
        E0 = Ellipsoid(np.zeros(2), np.diag([4.0, 1.0]))
        E = (R @ E0) + b

        l = np.array([1.0, 0.0])
        val, _ = E.rho(l)

        val0, _ = E0.rho(R.T @ l)
        expected = l @ b + val0
        assert val == pytest.approx(expected, abs=1e-10)

    def test_project_then_support(self):
        """Project a 3D ellipsoid to 2D, then check support.

        The support function of the projection in direction l2 should
        equal the support of the original in direction B l2.
        """
        E3 = Ellipsoid([1, 2, 3], np.diag([4.0, 9.0, 16.0]))
        B = np.array([[1, 0], [0, 1], [0, 0]], dtype=float)
        E2 = E3.projection(B)

        l2 = np.array([1.0, 0.0])
        val_proj, _ = E2.rho(l2)

        l3 = B @ l2
        val_orig, _ = E3.rho(l3)
        assert val_proj == pytest.approx(val_orig, abs=1e-10)

    def test_inv_polar_identity_for_unit_ball(self):
        """For the unit ball at origin:
        inv(E(0, I)) = E(0, I), polar(E(0, I)) = E(0, I).

        Both operations are identity for the unit ball.
        """
        E = Ellipsoid(np.zeros(3), np.eye(3))
        assert E.inv() == E
        assert E.polar() == E

    def test_linear_map_preserves_support_identity(self):
        """rho(l | A@E) = rho(A^T l | E) — the key support function
        identity for linear maps.
        """
        A = np.array([[2, 1], [0, 3]])
        E = Ellipsoid([1, -1], np.array([[3.0, 0.5], [0.5, 2.0]]))
        AE = A @ E

        rng = np.random.default_rng(7)
        for _ in range(20):
            l = rng.standard_normal(2)
            val_mapped, _ = AE.rho(l)
            val_original, _ = E.rho(A.T @ l)
            assert val_mapped == pytest.approx(val_original, abs=1e-10)


# ======================================================================
# Main entry point
# ======================================================================

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