"""Phase 3D tests — Internal reach set approximation.

Tests the internal shape matrix ODE (with and without disturbance)
and the reach_ia() function.  Key property: the IA must always be
contained inside the EA for the same direction.
"""

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, reach_ia, ReachSet


# =====================================================================
# ReachSet IA accessors
# =====================================================================

class TestReachSetIA:
    def test_ia_ellipsoid(self):
        rs = ReachSet()
        rs.t_eval = np.array([0.0, 1.0])
        rs.centers = np.array([[0, 0], [1, 0]])
        rs.ia_shapes = np.zeros((2, 1, 2, 2))
        rs.ia_shapes[0, 0] = np.eye(2)
        rs.ia_shapes[1, 0] = 0.5 * np.eye(2)
        rs.n_dirs = 1
        rs.n = 2

        E0 = rs.ia_ellipsoid(0, 0)
        assert isinstance(E0, Ellipsoid)
        np.testing.assert_array_equal(E0.center, [0, 0])

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

    def test_ia_ellipsoids_at(self):
        rs = ReachSet()
        rs.t_eval = np.array([0.0])
        rs.centers = np.array([[0, 0]])
        rs.ia_shapes = np.zeros((1, 3, 2, 2))
        for j in range(3):
            rs.ia_shapes[0, j] = (j + 1) * np.eye(2)
        rs.n_dirs = 3
        rs.n = 2

        ells = rs.ia_ellipsoids_at(0)
        assert len(ells) == 3

    def test_no_ia_raises(self):
        rs = ReachSet()
        with pytest.raises(ValueError, match="No internal"):
            rs.ia_ellipsoid(0, 0)


# =====================================================================
# reach_ia — smoke tests
# =====================================================================

class TestReachIASmoke:
    def test_returns_reach_set(self):
        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_ia(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.ia_shapes.shape == (20, 2, 2, 2)
        assert rs.n_dirs == 2
        assert rs.n == 2

    def test_initial_shape_is_Q0(self):
        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_ia(sys, X0, L0, (0, 1), N_steps=10)

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

    def test_1d_system(self):
        """1D system: IA should also grow over time with control input."""
        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_ia(sys, X0, L0, (0, 2), N_steps=50)

        q0_val = rs.ia_shapes[0, 0, 0, 0]
        qf_val = rs.ia_shapes[-1, 0, 0, 0]
        assert qf_val > q0_val

    def test_empty_system_raises(self):
        with pytest.raises(ValueError, match="non-empty"):
            reach_ia(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_ia(sys, Ellipsoid(), np.eye(2), (0, 1))


# =====================================================================
# IA ⊆ EA — the fundamental sandwich property
# =====================================================================

class TestIAInsideEA:
    """The internal approximation must lie inside the external approximation
    for the same direction at every time step.  We check this via the
    support function: ρ(l | IA) ≤ ρ(l | EA) for all l."""

    @staticmethod
    def _check_ia_inside_ea(sys, X0, L0, t_span, N_steps=30):
        rs_ea = reach_ea(sys, X0, L0, t_span, N_steps=N_steps)
        rs_ia = reach_ia(sys, X0, L0, t_span, N_steps=N_steps)

        n = rs_ea.n
        for k in range(N_steps):
            for j in range(rs_ea.n_dirs):
                E_ea = rs_ea.ea_ellipsoid(k, j)
                E_ia = rs_ia.ia_ellipsoid(k, j)

                # Check via eigenvalues of Q_ia vs Q_ea:
                # The IA shape should generally have smaller or equal eigenvalues.
                # More precisely: ρ(l | IA) ≤ ρ(l | EA) for the tightening direction.
                l = L0[:, j]
                rho_ea, _ = E_ea.rho(l)
                rho_ia, _ = E_ia.rho(l)

                assert float(rho_ia) <= float(rho_ea) + 1e-4, (
                    f"IA pokes outside EA at t={rs_ea.t_eval[k]:.3f}, "
                    f"dir={j}: ρ_ia={float(rho_ia):.6f} > ρ_ea={float(rho_ea):.6f}"
                )

    def test_stable_system_sandwich(self):
        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)
        self._check_ia_inside_ea(sys, X0, L0, (0, 2))

    def test_integrator_sandwich(self):
        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)
        self._check_ia_inside_ea(sys, X0, L0, (0, 1))

    def test_zero_A_sandwich(self):
        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]])
        self._check_ia_inside_ea(sys, X0, L0, (0, 1))


# =====================================================================
# reach_ia — known solutions
# =====================================================================

class TestReachIAKnownSolutions:
    def test_stable_system_shape_bounded(self):
        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_ia(sys, X0, L0, (0, 5), N_steps=50)

        max_eig = max(np.max(np.linalg.eigvalsh(rs.ia_shapes[k, 0]))
                      for k in range(50))
        assert max_eig < 100

    def test_ia_grows_with_control(self):
        """With A=0 and control, IA should grow monotonically."""
        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), 0.5 * np.eye(2))
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ia(sys, X0, L0, (0, 1), N_steps=15)

        eigs = [np.max(np.linalg.eigvalsh(rs.ia_shapes[k, 0]))
                for k in range(15)]
        for i in range(1, len(eigs)):
            assert eigs[i] >= eigs[i - 1] - ABS_TOL

    def test_multiple_directions_differ(self):
        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_ia(sys, X0, L0, (0, 1), N_steps=20)

        Q0 = rs.ia_shapes[-1, 0]
        Q1 = rs.ia_shapes[-1, 1]
        assert not np.allclose(Q0, Q1, atol=0.1)


# =====================================================================
# reach_ia with disturbance
# =====================================================================

class TestReachIADisturbance:
    def test_disturbance_shrinks_guaranteed_ia(self):
        """Adversarial disturbance shrinks the IA (same as EA)."""
        A = -np.eye(2)
        B = np.eye(2)
        U = Ellipsoid(np.zeros(2), np.eye(2))

        sys_no = LinSys(A, B, U)
        G = np.eye(2)
        V = Ellipsoid(np.zeros(2), 0.3 * np.eye(2))
        sys_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_ia(sys_no, X0, L0, (0, 2), N_steps=25)
        rs_dist = reach_ia(sys_dist, X0, L0, (0, 2), N_steps=25)

        eig_no = np.max(np.linalg.eigvalsh(rs_no.ia_shapes[-1, 0]))
        eig_dist = np.max(np.linalg.eigvalsh(rs_dist.ia_shapes[-1, 0]))
        assert eig_dist <= eig_no + ABS_TOL

    def test_disturbance_shapes_are_psd(self):
        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_ia(sys, X0, L0, (0, 2), N_steps=20)

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

    def test_ia_inside_ea_with_disturbance(self):
        """Sandwich property should hold even with disturbance."""
        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.0]])

        rs_ea = reach_ea(sys, X0, L0, (0, 2), N_steps=20)
        rs_ia = reach_ia(sys, X0, L0, (0, 2), N_steps=20)

        for k in range(20):
            l = L0[:, 0]
            rho_ea = float(rs_ea.ea_ellipsoid(k, 0).rho(l)[0])
            rho_ia = float(rs_ia.ia_ellipsoid(k, 0).rho(l)[0])
            assert rho_ia <= rho_ea + 1e-4


# =====================================================================
# Property tests
# =====================================================================

class TestReachIAProperties:
    def test_shapes_are_symmetric(self):
        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_ia(sys, X0, L0, (0, 2), N_steps=20)

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

    def test_center_matches_ea(self):
        """IA and EA should have the same center trajectory."""
        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_ea = reach_ea(sys, X0, L0, (0, 2), N_steps=20)
        rs_ia = reach_ia(sys, X0, L0, (0, 2), N_steps=20)

        np.testing.assert_allclose(rs_ia.centers, rs_ea.centers, atol=1e-6)

    def test_3d_system(self):
        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_ia(sys, X0, L0, (0, 1), N_steps=10)

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

    def test_single_step(self):
        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_ia(sys, X0, L0, (0, 1), N_steps=2)
        assert rs.ia_shapes.shape[0] == 2
