"""Phase 3E tests — Discrete-time reach set approximation.

Tests reach_ea_d and reach_ia_d for discrete-time systems:
    x[k+1] = A x[k] + B u[k]  (+ G v[k])

Key properties checked:
    - IA ⊆ EA (sandwich) for the same direction
    - Shapes are symmetric PSD
    - Centers evolve correctly
    - Known analytical solutions
"""

import numpy as np
import pytest

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_d, reach_ia_d, ReachSet


# =====================================================================
# Helper: build a simple discrete-time system
# =====================================================================

def _simple_sys(A, B=None, U=None, G=None, V=None):
    """Build a discrete LinSys with default B=I, U=unit ball."""
    n = A.shape[0]
    if B is None:
        B = np.eye(n)
    if U is None:
        U = Ellipsoid(np.zeros(B.shape[1]), np.eye(B.shape[1]))
    return LinSys(A, B, U, G, V, discrete=True)


# =====================================================================
# reach_ea_d — smoke tests
# =====================================================================

class TestReachEADSmoke:
    def test_returns_reach_set(self):
        sys = _simple_sys(0.9 * np.eye(2))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.eye(2)
        rs = reach_ea_d(sys, X0, L0, N_steps=5)

        assert isinstance(rs, ReachSet)
        assert rs.t_eval.shape == (6,)       # k=0..5
        assert rs.centers.shape == (6, 2)
        assert rs.ea_shapes.shape == (6, 2, 2, 2)
        assert rs.n_dirs == 2

    def test_initial_shape_is_Q0(self):
        Q0 = 4 * np.eye(2)
        sys = _simple_sys(np.eye(2))
        X0 = Ellipsoid(np.zeros(2), Q0)
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ea_d(sys, X0, L0, N_steps=3)
        np.testing.assert_allclose(rs.ea_shapes[0, 0], Q0, atol=1e-10)

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

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


# =====================================================================
# reach_ia_d — smoke tests
# =====================================================================

class TestReachIADSmoke:
    def test_returns_reach_set(self):
        sys = _simple_sys(0.9 * np.eye(2))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.eye(2)
        rs = reach_ia_d(sys, X0, L0, N_steps=5)

        assert isinstance(rs, ReachSet)
        assert rs.ia_shapes.shape == (6, 2, 2, 2)

    def test_initial_shape_is_Q0(self):
        Q0 = 3 * np.eye(2)
        sys = _simple_sys(np.eye(2))
        X0 = Ellipsoid(np.zeros(2), Q0)
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ia_d(sys, X0, L0, N_steps=3)
        np.testing.assert_allclose(rs.ia_shapes[0, 0], Q0, atol=1e-10)

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


# =====================================================================
# Center trajectory
# =====================================================================

class TestCenters:
    def test_center_with_identity_dynamics(self):
        """A=I, B=I, u=0: center should stay at origin."""
        sys = _simple_sys(np.eye(2))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ea_d(sys, X0, L0, N_steps=5)
        for k in range(6):
            np.testing.assert_allclose(rs.centers[k], [0, 0], atol=1e-12)

    def test_center_with_shift(self):
        """A=I, B=I, U centered at [1,0]: center moves linearly."""
        U = Ellipsoid([1.0, 0.0], np.eye(2))
        sys = LinSys(np.eye(2), np.eye(2), U, discrete=True)
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ea_d(sys, X0, L0, N_steps=5)

        for k in range(6):
            np.testing.assert_allclose(rs.centers[k], [k, 0], atol=1e-10)

    def test_center_with_scaling(self):
        """A = 0.5*I, B=0, no control: center decays by 0.5 each step."""
        sys = LinSys(0.5 * np.eye(2), discrete=True)
        X0 = Ellipsoid([4.0, 2.0], np.eye(2))
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ea_d(sys, X0, L0, N_steps=3)

        np.testing.assert_allclose(rs.centers[0], [4, 2], atol=1e-12)
        np.testing.assert_allclose(rs.centers[1], [2, 1], atol=1e-12)
        np.testing.assert_allclose(rs.centers[2], [1, 0.5], atol=1e-12)

    def test_ia_and_ea_same_centers(self):
        """IA and EA should produce identical center trajectories."""
        sys = _simple_sys(0.9 * np.eye(2))
        X0 = Ellipsoid([1.0, 2.0], np.eye(2))
        L0 = np.eye(2)
        rs_ea = reach_ea_d(sys, X0, L0, N_steps=5)
        rs_ia = reach_ia_d(sys, X0, L0, N_steps=5)
        np.testing.assert_allclose(rs_ia.centers, rs_ea.centers, atol=1e-12)


# =====================================================================
# Shape properties
# =====================================================================

class TestShapeProperties:
    def test_ea_shapes_symmetric(self):
        sys = _simple_sys(0.9 * np.eye(2))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.eye(2)
        rs = reach_ea_d(sys, X0, L0, N_steps=10)
        for k in range(11):
            for j in range(2):
                Q = rs.ea_shapes[k, j]
                np.testing.assert_allclose(Q, Q.T, atol=1e-10)

    def test_ea_shapes_psd(self):
        sys = _simple_sys(0.9 * np.eye(2))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.eye(2)
        rs = reach_ea_d(sys, X0, L0, N_steps=10)
        for k in range(11):
            for j in range(2):
                eigvals = np.linalg.eigvalsh(rs.ea_shapes[k, j])
                assert np.min(eigvals) >= -ABS_TOL

    def test_ia_shapes_symmetric(self):
        sys = _simple_sys(0.9 * np.eye(2))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.eye(2)
        rs = reach_ia_d(sys, X0, L0, N_steps=10)
        for k in range(11):
            for j in range(2):
                Q = rs.ia_shapes[k, j]
                np.testing.assert_allclose(Q, Q.T, atol=1e-10)

    def test_ia_shapes_psd(self):
        sys = _simple_sys(0.9 * np.eye(2))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.eye(2)
        rs = reach_ia_d(sys, X0, L0, N_steps=10)
        for k in range(11):
            for j in range(2):
                eigvals = np.linalg.eigvalsh(rs.ia_shapes[k, j])
                assert np.min(eigvals) >= -ABS_TOL


# =====================================================================
# IA ⊆ EA (sandwich property)
# =====================================================================

class TestSandwich:
    @staticmethod
    def _check_sandwich(sys, X0, L0, N_steps):
        rs_ea = reach_ea_d(sys, X0, L0, N_steps=N_steps)
        rs_ia = reach_ia_d(sys, X0, L0, N_steps=N_steps)

        for k in range(N_steps + 1):
            for j in range(L0.shape[1]):
                E_ea = rs_ea.ea_ellipsoid(k, j)
                E_ia = rs_ia.ia_ellipsoid(k, j)
                l = L0[:, j]
                rho_ea = float(E_ea.rho(l)[0])
                rho_ia = float(E_ia.rho(l)[0])
                assert rho_ia <= rho_ea + 1e-4, (
                    f"Sandwich violated at k={k}, dir={j}: "
                    f"ρ_ia={rho_ia:.6f} > ρ_ea={rho_ea:.6f}"
                )

    def test_stable_sandwich(self):
        sys = _simple_sys(0.9 * np.eye(2))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1, 0], [0, 1]], dtype=float)
        self._check_sandwich(sys, X0, L0, 8)

    def test_identity_sandwich(self):
        sys = _simple_sys(np.eye(2))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1.0], [0.0]])
        self._check_sandwich(sys, X0, L0, 5)

    def test_rotation_sandwich(self):
        th = np.pi / 6
        A = np.array([[np.cos(th), -np.sin(th)],
                       [np.sin(th),  np.cos(th)]])
        sys = _simple_sys(A)
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.eye(2)
        self._check_sandwich(sys, X0, L0, 6)


# =====================================================================
# Known solutions
# =====================================================================

class TestKnownSolutions:
    def test_identity_grows_monotonically(self):
        """A=I with control: EA shapes should grow every step."""
        sys = _simple_sys(np.eye(2))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ea_d(sys, X0, L0, N_steps=5)

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

    def test_contractive_bounded(self):
        """Strongly stable system: shapes stay bounded."""
        sys = _simple_sys(0.5 * np.eye(2))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.eye(2)
        rs = reach_ea_d(sys, X0, L0, N_steps=50)

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

    def test_1d_identity_ea(self):
        """1D: A=1, B=1, U=[-1,1]. EA at step k = (1 + k)^2 * Q0_eff."""
        A = np.array([[1.0]])
        B = np.array([[1.0]])
        U = Ellipsoid([0.0], np.array([[1.0]]))
        sys = LinSys(A, B, U, discrete=True)
        X0 = Ellipsoid([0.0], np.array([[1.0]]))
        L0 = np.array([[1.0]])
        rs = reach_ea_d(sys, X0, L0, N_steps=3)

        # EA shapes should be monotonically increasing (1D scalar)
        vals = [rs.ea_shapes[k, 0, 0, 0] for k in range(4)]
        for i in range(1, 4):
            assert vals[i] > vals[i - 1]

    def test_multiple_directions_differ(self):
        """Non-isotropic system should give different shapes per direction."""
        A = np.array([[1.0, 0.5], [0.0, 0.8]])
        sys = _simple_sys(A)
        X0 = Ellipsoid(np.zeros(2), 0.1 * np.eye(2))
        L0 = np.array([[1, 0], [0, 1]], dtype=float)
        rs = reach_ea_d(sys, X0, L0, N_steps=5)

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


# =====================================================================
# Disturbance
# =====================================================================

class TestDisturbance:
    def test_disturbance_shrinks_ea(self):
        """Adversarial disturbance shrinks the guaranteed EA."""
        A = 0.9 * np.eye(2)
        B = np.eye(2)
        U = Ellipsoid(np.zeros(2), np.eye(2))

        sys_no = LinSys(A, B, U, discrete=True)
        G = np.eye(2)
        V = Ellipsoid(np.zeros(2), 0.3 * np.eye(2))
        sys_dist = LinSys(A, B, U, G, V, discrete=True)

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

        rs_no = reach_ea_d(sys_no, X0, L0, N_steps=8)
        rs_dist = reach_ea_d(sys_dist, X0, L0, N_steps=8)

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

    def test_disturbance_shapes_psd(self):
        A = 0.9 * 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, discrete=True)
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.eye(2)

        rs = reach_ea_d(sys, X0, L0, N_steps=5)
        for k in range(6):
            for j in range(2):
                eigvals = np.linalg.eigvalsh(rs.ea_shapes[k, j])
                assert np.min(eigvals) >= -ABS_TOL

    def test_disturbance_sandwich(self):
        A = 0.9 * 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, discrete=True)
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1.0], [0.0]])

        rs_ea = reach_ea_d(sys, X0, L0, N_steps=5)
        rs_ia = reach_ia_d(sys, X0, L0, N_steps=5)

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


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

class TestEdgeCases:
    def test_single_step(self):
        sys = _simple_sys(np.eye(2))
        X0 = Ellipsoid(np.zeros(2), np.eye(2))
        L0 = np.array([[1.0], [0.0]])
        rs = reach_ea_d(sys, X0, L0, N_steps=1)
        assert rs.ea_shapes.shape[0] == 2  # k=0 and k=1

    def test_3d_system(self):
        sys = _simple_sys(0.9 * np.eye(3))
        X0 = Ellipsoid(np.zeros(3), np.eye(3))
        L0 = np.eye(3)
        rs = reach_ea_d(sys, X0, L0, N_steps=3)
        assert rs.ea_shapes.shape == (4, 3, 3, 3)

    def test_1d_system(self):
        A = np.array([[0.8]])
        B = np.array([[1.0]])
        U = Ellipsoid([0.0], np.array([[1.0]]))
        sys = LinSys(A, B, U, discrete=True)
        X0 = Ellipsoid([0.0], np.array([[1.0]]))
        L0 = np.array([[1.0]])
        rs_ea = reach_ea_d(sys, X0, L0, N_steps=5)
        rs_ia = reach_ia_d(sys, X0, L0, N_steps=5)
        assert rs_ea.ea_shapes.shape == (6, 1, 1, 1)
        assert rs_ia.ia_shapes.shape == (6, 1, 1, 1)
