"""Phase 3C tests — External reach set approximation.

Tests the external shape matrix ODE (with and without disturbance)
and the reach_ea() function against known analytical results.
"""

import numpy as np
import pytest
import scipy.linalg

from ellipsoidal_toolbox.core.config import ABS_TOL
from ellipsoidal_toolbox.core.ellipsoid import Ellipsoid
from ellipsoidal_toolbox.systems.linsys import LinSys
from ellipsoidal_toolbox.systems.reach import reach_ea, ReachSet


# =====================================================================
# ReachSet data structure
# =====================================================================

class TestReachSet:
    def test_ea_ellipsoid(self):
        """Extract a single EA ellipsoid from the reach set."""
        rs = ReachSet()
        rs.t_eval = np.array([0.0, 1.0])
        rs.centers = np.array([[0, 0], [1, 0]])
        rs.ea_shapes = np.zeros((2, 1, 2, 2))
        rs.ea_shapes[0, 0] = np.eye(2)
        rs.ea_shapes[1, 0] = 2 * np.eye(2)
        rs.n_dirs = 1
        rs.n = 2

        E0 = rs.ea_ellipsoid(0, 0)
        assert isinstance(E0, Ellipsoid)
        np.testing.assert_array_equal(E0.center, [0, 0])
        np.testing.assert_array_equal(E0.shape, np.eye(2))

        E1 = rs.ea_ellipsoid(1, 0)
        np.testing.assert_array_equal(E1.center, [1, 0])

    def test_ea_ellipsoids_at(self):
        rs = ReachSet()
        rs.t_eval = np.array([0.0])
        rs.centers = np.array([[0, 0]])
        rs.ea_shapes = np.zeros((1, 2, 2, 2))
        rs.ea_shapes[0, 0] = np.eye(2)
        rs.ea_shapes[0, 1] = 2 * np.eye(2)
        rs.n_dirs = 2
        rs.n = 2

        ells = rs.ea_ellipsoids_at(0)
        assert len(ells) == 2

    def test_no_ea_raises(self):
        rs = ReachSet()
        with pytest.raises(ValueError, match="No external"):
            rs.ea_ellipsoid(0, 0)


# =====================================================================
# reach_ea — basic smoke tests
# =====================================================================

class TestReachEASmoke:
    def test_returns_reach_set(self):
        """reach_ea returns a ReachSet with correct dimensions."""
        A = np.array([[-1.0, 0], [0, -1.0]])
        B = np.eye(2)
        U = Ellipsoid(np.zeros(2), np.eye(2))
        sys = LinSys(A, B, U)
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.eye(2)
        rs = reach_ea(sys, X0, L0, (0, 1), N_steps=20)

        assert isinstance(rs, ReachSet)
        assert rs.t_eval.shape == (20,)
        assert rs.centers.shape == (20, 2)
        assert rs.ea_shapes.shape == (20, 2, 2, 2)
        assert rs.n_dirs == 2
        assert rs.n == 2

    def test_initial_shape_is_Q0(self):
        """At t=0, the EA shape should be Q0."""
        Q0 = 4 * np.eye(2)
        A = np.zeros((2, 2))
        B = np.eye(2)
        U = Ellipsoid(np.zeros(2), np.eye(2))
        sys = LinSys(A, B, U)
        X0 = Ellipsoid(np.zeros(2), Q0)
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ea(sys, X0, L0, (0, 1), N_steps=10)

        np.testing.assert_allclose(rs.ea_shapes[0, 0], Q0, atol=1e-6)

    def test_1d_system(self):
        """1D system: dx/dt = 0*x + u, |u| <= 1."""
        A = np.array([[0.0]])
        B = np.array([[1.0]])
        U = Ellipsoid([0.0], np.array([[1.0]]))
        sys = LinSys(A, B, U)
        X0 = Ellipsoid([0.0], np.array([[0.01]]))
        L0 = np.array([[1.0]])
        rs = reach_ea(sys, X0, L0, (0, 2), N_steps=50)

        # The set grows linearly; shape should be increasing
        q0 = rs.ea_shapes[0, 0, 0, 0]
        qf = rs.ea_shapes[-1, 0, 0, 0]
        assert qf > q0

    def test_empty_system_raises(self):
        with pytest.raises(ValueError, match="non-empty"):
            reach_ea(LinSys(), Ellipsoid(np.zeros(2), np.eye(2)),
                     np.eye(2), (0, 1))

    def test_empty_X0_raises(self):
        sys = LinSys(np.eye(2))
        with pytest.raises(ValueError, match="non-empty"):
            reach_ea(sys, Ellipsoid(), np.eye(2), (0, 1))


# =====================================================================
# reach_ea — known solutions
# =====================================================================

class TestReachEAKnownSolutions:
    def test_stable_system_shape_bounded(self):
        """Stable system (A = -I): reach set should not blow up."""
        A = -np.eye(2)
        B = np.eye(2)
        U = Ellipsoid(np.zeros(2), np.eye(2))
        sys = LinSys(A, B, U)
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1, 0], [0, 1]], dtype=float)
        rs = reach_ea(sys, X0, L0, (0, 5), N_steps=50)

        # Shape matrices should stay bounded
        max_eig = max(np.max(np.linalg.eigvalsh(rs.ea_shapes[k, 0]))
                      for k in range(50))
        assert max_eig < 100  # generous bound

    def test_autonomous_system_shrinks(self):
        """Autonomous stable system (no input): ellipsoid shrinks via dynamics."""
        A = np.array([[-2.0, 0], [0, -2.0]])
        sys = LinSys(A)
        Q0 = 4 * np.eye(2)
        X0 = Ellipsoid(np.zeros(2), Q0)
        L0 = np.array([[1.0], [0.0]])

        # Without control input, BPB=0.
        # The ODE becomes dQ/dt = AQ + QA' + (0/p2)*Q + (p2/0)*0
        # With p1 → 0, the extra terms vanish (both clamped to ABS_TOL).
        # So dQ/dt ≈ AQ + QA' = -4Q, Q(t) ≈ Q0 * exp(-4t)
        rs = reach_ea(sys, X0, L0, (0, 2), N_steps=30)

        # Final shape should be much smaller than initial
        q0_eig = np.max(np.linalg.eigvalsh(rs.ea_shapes[0, 0]))
        qf_eig = np.max(np.linalg.eigvalsh(rs.ea_shapes[-1, 0]))
        assert qf_eig < q0_eig

    def test_integrator_with_unit_ball(self):
        """Double integrator with unit ball control.

        dx1/dt = x2
        dx2/dt = u, |u| <= 1

        The reachable set grows over time.
        """
        A = np.array([[0, 1], [0, 0]])
        B = np.array([[0], [1]])
        U = Ellipsoid([0.0], np.array([[1.0]]))
        sys = LinSys(A, B, U)
        X0 = Ellipsoid(np.zeros(2), 0.01 * np.eye(2))
        L0 = np.array([[1, 0], [0, 1]], dtype=float)
        rs = reach_ea(sys, X0, L0, (0, 2), N_steps=40)

        # Both directions should show growth
        vol_0 = np.linalg.det(rs.ea_shapes[0, 0])
        vol_f = np.linalg.det(rs.ea_shapes[-1, 0])
        assert vol_f > vol_0

    def test_multiple_directions_give_different_shapes(self):
        """Different directions produce different shape matrices."""
        A = np.array([[0, 1], [0, 0]])
        B = np.array([[0], [1]])
        U = Ellipsoid([0.0], np.array([[1.0]]))
        sys = LinSys(A, B, U)
        X0 = Ellipsoid(np.zeros(2), 0.01 * np.eye(2))
        L0 = np.array([[1, 0], [0, 1]], dtype=float)
        rs = reach_ea(sys, X0, L0, (0, 1), N_steps=20)

        # At final time, shapes for the two directions should differ
        # (different approximation tightness)
        Q_dir0 = rs.ea_shapes[-1, 0]
        Q_dir1 = rs.ea_shapes[-1, 1]
        assert not np.allclose(Q_dir0, Q_dir1, atol=0.1)


# =====================================================================
# reach_ea with disturbance
# =====================================================================

class TestReachEADisturbance:
    def test_disturbance_shrinks_guaranteed_set(self):
        """Adversarial disturbance shrinks the guaranteed reach set.

        The external reach set ODE uses Minkowski difference for
        disturbance (Kurzhanski-Valyev framework), so the guaranteed
        reachable set is smaller when disturbance is present.
        """
        A = -np.eye(2)
        B = np.eye(2)
        U = Ellipsoid(np.zeros(2), np.eye(2))

        sys_no_dist = LinSys(A, B, U)
        G = np.eye(2)
        V = Ellipsoid(np.zeros(2), 0.5 * np.eye(2))
        sys_with_dist = LinSys(A, B, U, G, V)

        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1.0], [0.0]])

        rs_no = reach_ea(sys_no_dist, X0, L0, (0, 3), N_steps=30)
        rs_yes = reach_ea(sys_with_dist, X0, L0, (0, 3), N_steps=30)

        eig_no = np.max(np.linalg.eigvalsh(rs_no.ea_shapes[-1, 0]))
        eig_yes = np.max(np.linalg.eigvalsh(rs_yes.ea_shapes[-1, 0]))
        assert eig_yes <= eig_no + ABS_TOL

    def test_disturbance_shapes_are_psd(self):
        """All shape matrices should be PSD."""
        A = -np.eye(2)
        B = np.eye(2)
        U = Ellipsoid(np.zeros(2), np.eye(2))
        G = np.eye(2)
        V = Ellipsoid(np.zeros(2), 0.1 * np.eye(2))
        sys = LinSys(A, B, U, G, V)
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1, 0], [0, 1]], dtype=float)

        rs = reach_ea(sys, X0, L0, (0, 2), N_steps=20)

        for k in range(20):
            for j in range(2):
                eigvals = np.linalg.eigvalsh(rs.ea_shapes[k, j])
                assert np.min(eigvals) >= -ABS_TOL, (
                    f"Negative eigenvalue at t={rs.t_eval[k]}, dir={j}: "
                    f"{np.min(eigvals)}"
                )


# =====================================================================
# reach_ea — property tests
# =====================================================================

class TestReachEAProperties:
    def test_shapes_are_symmetric(self):
        """All EA shape matrices should be symmetric."""
        A = np.array([[0, 1], [-0.5, -1]])
        B = np.array([[0], [1]])
        U = Ellipsoid([0.0], np.array([[1.0]]))
        sys = LinSys(A, B, U)
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ea(sys, X0, L0, (0, 2), N_steps=20)

        for k in range(20):
            Q = rs.ea_shapes[k, 0]
            np.testing.assert_allclose(Q, Q.T, atol=1e-10)

    def test_center_trajectory_matches_standalone(self):
        """Centers from reach_ea should match compute_center."""
        from ellipsoidal_toolbox.systems.ode_utils import compute_center

        A = -np.eye(2)
        B = np.eye(2)
        U = Ellipsoid([1.0, 0.5], np.eye(2))
        sys = LinSys(A, B, U)
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1.0], [0.0]])

        rs = reach_ea(sys, X0, L0, (0, 2), N_steps=20)
        centers_standalone = compute_center(sys, np.zeros(2), (0, 2), rs.t_eval)
        np.testing.assert_allclose(rs.centers, centers_standalone, atol=1e-6)

    def test_3d_system(self):
        """3D system should work without errors."""
        A = -np.eye(3)
        B = np.eye(3)
        U = Ellipsoid(np.zeros(3), np.eye(3))
        sys = LinSys(A, B, U)
        X0 = Ellipsoid(np.zeros(3), np.eye(3))
        L0 = np.eye(3)
        rs = reach_ea(sys, X0, L0, (0, 1), N_steps=10)

        assert rs.ea_shapes.shape == (10, 3, 3, 3)
        assert rs.n_dirs == 3


# =====================================================================
# Edge cases
# =====================================================================

class TestEdgeCases:
    def test_zero_A_constant_growth(self):
        """A=0: reach set grows from input only."""
        A = np.zeros((2, 2))
        B = np.eye(2)
        U = Ellipsoid(np.zeros(2), np.eye(2))
        sys = LinSys(A, B, U)
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ea(sys, X0, L0, (0, 1), N_steps=10)

        # Should grow monotonically
        eigs = [np.max(np.linalg.eigvalsh(rs.ea_shapes[k, 0]))
                for k in range(10)]
        for i in range(1, len(eigs)):
            assert eigs[i] >= eigs[i - 1] - ABS_TOL

    def test_single_direction(self):
        """Single direction column should work."""
        sys = LinSys(-np.eye(2), np.eye(2),
                      Ellipsoid(np.zeros(2), np.eye(2)))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ea(sys, X0, L0, (0, 1), N_steps=5)
        assert rs.n_dirs == 1

    def test_single_step(self):
        """N_steps=2 (minimum for a span) should work."""
        sys = LinSys(-np.eye(2), np.eye(2),
                      Ellipsoid(np.zeros(2), np.eye(2)))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ea(sys, X0, L0, (0, 1), N_steps=2)
        assert rs.ea_shapes.shape[0] == 2
