"""Phase 3B tests — ODE infrastructure.

Tests state transition matrix, center trajectory, direction evolution,
and precomputation helpers for both LTI and LTV systems.
"""

import numpy as np
import pytest
import scipy.linalg

from ellipsoidal_toolbox.core.ellipsoid import Ellipsoid
from ellipsoidal_toolbox.systems.linsys import LinSys
from ellipsoidal_toolbox.systems.ode_utils import (
    ode_solve,
    compute_stm,
    compute_stm_inv,
    compute_center,
    evolve_directions,
    precompute_bounds,
    precompute_all,
)


# =====================================================================
# ODE solver wrapper
# =====================================================================

class TestOdeSolve:
    def test_simple_decay(self):
        """dx/dt = -x, x(0)=1 → x(t) = exp(-t)."""
        sol = ode_solve(lambda t, x: -x, (0, 2), np.array([1.0]),
                        t_eval=np.array([0.0, 1.0, 2.0]))
        assert sol.status == 0
        np.testing.assert_allclose(sol.y[0], np.exp(-sol.t), rtol=1e-5)

    def test_oscillator(self):
        """2D harmonic oscillator: known analytical solution."""
        def rhs(t, x):
            return np.array([x[1], -x[0]])

        sol = ode_solve(rhs, (0, np.pi), np.array([1.0, 0.0]),
                        t_eval=np.array([0.0, np.pi / 2, np.pi]))
        assert sol.status == 0
        np.testing.assert_allclose(sol.y[:, -1], [-1.0, 0.0], atol=1e-4)


# =====================================================================
# State transition matrix — LTI
# =====================================================================

class TestSTM_LTI:
    def test_identity_at_t0(self):
        """Φ(t0, t0) = I."""
        sys = LinSys(np.array([[0, 1], [-1, 0]]))
        t_eval = np.array([0.0])
        Phi = compute_stm(sys, (0, 1), t_eval)
        np.testing.assert_allclose(Phi[0], np.eye(2), atol=1e-10)

    def test_scalar_decay(self):
        """A = [[-1]], Φ(t) = exp(-t)."""
        sys = LinSys(np.array([[-1.0]]))
        t_eval = np.linspace(0, 2, 5)
        Phi = compute_stm(sys, (0, 2), t_eval)
        for k, t in enumerate(t_eval):
            expected = np.exp(-t)
            assert Phi[k, 0, 0] == pytest.approx(expected, rel=1e-5)

    def test_rotation(self):
        """A = [[0, 1], [-1, 0]] gives rotation matrices."""
        A = np.array([[0, 1], [-1, 0]])
        sys = LinSys(A)
        t_eval = np.array([0, np.pi / 2, np.pi])
        Phi = compute_stm(sys, (0, np.pi), t_eval)
        # At t=pi/2: rotation by 90 degrees
        expected = np.array([[0, 1], [-1, 0]])
        np.testing.assert_allclose(Phi[1], expected, atol=1e-5)
        # At t=pi: rotation by 180 degrees
        np.testing.assert_allclose(Phi[2], -np.eye(2), atol=1e-5)

    def test_expm_consistency(self):
        """Compare with scipy.linalg.expm directly."""
        A = np.array([[1, 0.5], [0, -1]])
        sys = LinSys(A)
        t_eval = np.array([0.0, 0.5, 1.0])
        Phi = compute_stm(sys, (0, 1), t_eval)
        for k, t in enumerate(t_eval):
            np.testing.assert_allclose(Phi[k], scipy.linalg.expm(A * t), atol=1e-10)

    def test_3d_system(self):
        """3D system STM at t=0 is identity."""
        A = np.array([[0, 1, 0], [0, 0, 1], [-1, -1, -1]])
        sys = LinSys(A)
        t_eval = np.array([0.0, 0.5])
        Phi = compute_stm(sys, (0, 0.5), t_eval)
        np.testing.assert_allclose(Phi[0], np.eye(3), atol=1e-10)
        np.testing.assert_allclose(Phi[1], scipy.linalg.expm(A * 0.5), atol=1e-8)


# =====================================================================
# State transition matrix — LTV
# =====================================================================

class TestSTM_LTV:
    def test_ltv_identity_at_t0(self):
        """LTV system: Φ(t0) = I."""
        def A_func(t):
            return np.array([[0, 1 + 0.1 * t], [-1, 0]])

        sys = LinSys(A_func)
        t_eval = np.array([0.0])
        Phi = compute_stm(sys, (0, 1), t_eval)
        np.testing.assert_allclose(Phi[0], np.eye(2), atol=1e-10)

    def test_ltv_scalar(self):
        """A(t) = [[-t]], Φ(t) = exp(-t²/2)."""
        def A_func(t):
            return np.array([[-t]])

        sys = LinSys(A_func)
        t_eval = np.linspace(0, 2, 20)
        Phi = compute_stm(sys, (0, 2), t_eval)
        for k, t in enumerate(t_eval):
            expected = np.exp(-t ** 2 / 2)
            assert Phi[k, 0, 0] == pytest.approx(expected, rel=1e-3)

    def test_ltv_constant_A_matches_lti(self):
        """LTV with constant callable A should match LTI expm."""
        A = np.array([[0, 1], [-1, 0]])

        def A_func(t):
            return A

        sys_lti = LinSys(A)
        sys_ltv = LinSys(A_func)
        t_eval = np.linspace(0, 1, 10)

        Phi_lti = compute_stm(sys_lti, (0, 1), t_eval)
        Phi_ltv = compute_stm(sys_ltv, (0, 1), t_eval)
        np.testing.assert_allclose(Phi_ltv, Phi_lti, atol=1e-4)


# =====================================================================
# STM inverse
# =====================================================================

class TestSTMInverse:
    def test_inverse_product_is_identity(self):
        A = np.array([[1, 0.5], [0, -1]])
        sys = LinSys(A)
        t_eval = np.linspace(0, 1, 5)
        Phi = compute_stm(sys, (0, 1), t_eval)
        Phi_inv = compute_stm_inv(Phi)
        for k in range(len(t_eval)):
            np.testing.assert_allclose(Phi[k] @ Phi_inv[k], np.eye(2), atol=1e-10)


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

class TestCenterTrajectory:
    def test_autonomous_decay(self):
        """dx/dt = -x, x(0)=[1] → x(t) = exp(-t)."""
        sys = LinSys(np.array([[-1.0]]))
        t_eval = np.linspace(0, 2, 20)
        centers = compute_center(sys, np.array([1.0]), (0, 2), t_eval)
        expected = np.exp(-t_eval)
        np.testing.assert_allclose(centers[:, 0], expected, rtol=1e-4)

    def test_with_control_center(self):
        """dx/dt = -x + u, u center = [1], steady state → x = 1."""
        A = np.array([[-1.0]])
        B = np.array([[1.0]])
        U = Ellipsoid([1.0], np.array([[0.01]]))
        sys = LinSys(A, B, U)
        t_eval = np.linspace(0, 10, 50)
        centers = compute_center(sys, np.array([0.0]), (0, 10), t_eval)
        # Steady state: dx/dt = 0 → x = B*p / |a| = 1
        assert centers[-1, 0] == pytest.approx(1.0, abs=0.05)

    def test_with_disturbance_center(self):
        """dx/dt = x + Gq, G=I, q=[0.5, 0.5]."""
        A = np.array([[-1, 0], [0, -1]])
        G = np.eye(2)
        V = Ellipsoid([0.5, 0.5], 0.01 * np.eye(2))
        sys = LinSys(A, G=G, V=V)
        t_eval = np.linspace(0, 10, 50)
        centers = compute_center(sys, np.zeros(2), (0, 10), t_eval)
        # Steady state: x = -A^{-1} G q = [0.5, 0.5]
        np.testing.assert_allclose(centers[-1], [0.5, 0.5], atol=0.05)

    def test_2d_oscillator(self):
        """Harmonic oscillator: x(t) = [cos(t), -sin(t)]."""
        A = np.array([[0, 1], [-1, 0]])
        sys = LinSys(A)
        t_eval = np.array([0.0, np.pi / 2, np.pi])
        centers = compute_center(sys, np.array([1.0, 0.0]), (0, np.pi), t_eval)
        np.testing.assert_allclose(centers[1], [0.0, -1.0], atol=1e-4)
        np.testing.assert_allclose(centers[2], [-1.0, 0.0], atol=1e-4)

    def test_empty_system_raises(self):
        sys = LinSys()
        with pytest.raises(ValueError, match="non-empty"):
            compute_center(sys, np.array([0.0]), (0, 1), np.array([0.0]))


# =====================================================================
# Direction evolution
# =====================================================================

class TestDirectionEvolution:
    def test_identity_stm_preserves_directions(self):
        """Φ = I at all times → directions unchanged."""
        N = 5
        n = 2
        Phi = np.tile(np.eye(n), (N, 1, 1))
        L0 = np.array([[1, 0], [0, 1]], dtype=float)
        L = evolve_directions(Phi, L0, backward=True)
        for k in range(N):
            np.testing.assert_allclose(L[k], L0, atol=1e-10)

    def test_rotation_stm(self):
        """Rotation STM rotates direction vectors."""
        theta = np.pi / 4
        R = np.array([[np.cos(theta), -np.sin(theta)],
                       [np.sin(theta), np.cos(theta)]])
        Phi = np.array([np.eye(2), R])
        L0 = np.array([[1.0], [0.0]])
        L = evolve_directions(Phi, L0, backward=True)
        # backward: l = Phi^T l0
        expected = R.T @ L0
        np.testing.assert_allclose(L[1], expected, atol=1e-10)

    def test_multiple_directions(self):
        """Multiple direction columns evolve independently."""
        N = 3
        Phi = np.tile(np.eye(2), (N, 1, 1))
        L0 = np.array([[1, 0], [0, 1]], dtype=float)
        L = evolve_directions(Phi, L0, backward=False)
        for k in range(N):
            np.testing.assert_allclose(L[k], L0, atol=1e-10)


# =====================================================================
# Precompute bounds
# =====================================================================

class TestPrecomputeBounds:
    def test_bpb_identity(self):
        """B=I, U=unit ball → BPB = I, BPBsr = I."""
        A = np.eye(2)
        B = np.eye(2)
        U = Ellipsoid(np.zeros(2), np.eye(2))
        sys = LinSys(A, B, U)
        t_eval = np.array([0.0, 0.5, 1.0])
        BPB, BPBsr, GQG, GQGsr = precompute_bounds(sys, t_eval)
        for k in range(3):
            np.testing.assert_allclose(BPB[k], np.eye(2), atol=1e-10)
            np.testing.assert_allclose(BPBsr[k], np.eye(2), atol=1e-6)
        assert GQG is None
        assert GQGsr is None

    def test_scaled_control(self):
        """B=I, P=4*I → BPB = 4*I, BPBsr = 2*I."""
        sys = LinSys(np.eye(2), np.eye(2),
                      U=Ellipsoid(np.zeros(2), 4 * np.eye(2)))
        t_eval = np.array([0.0])
        BPB, BPBsr, _, _ = precompute_bounds(sys, t_eval)
        np.testing.assert_allclose(BPB[0], 4 * np.eye(2), atol=1e-10)
        np.testing.assert_allclose(BPBsr[0], 2 * np.eye(2), atol=1e-6)

    def test_with_disturbance(self):
        """System with disturbance computes GQG."""
        sys = LinSys(np.eye(2), G=np.eye(2),
                      V=Ellipsoid(np.zeros(2), 0.25 * np.eye(2)))
        t_eval = np.array([0.0])
        BPB, BPBsr, GQG, GQGsr = precompute_bounds(sys, t_eval)
        assert GQG is not None
        np.testing.assert_allclose(GQG[0], 0.25 * np.eye(2), atol=1e-10)
        np.testing.assert_allclose(GQGsr[0], 0.5 * np.eye(2), atol=1e-6)

    def test_no_control(self):
        """System without B → BPB = 0."""
        sys = LinSys(np.eye(2))
        t_eval = np.array([0.0])
        BPB, BPBsr, _, _ = precompute_bounds(sys, t_eval)
        np.testing.assert_allclose(BPB[0], np.zeros((2, 2)), atol=1e-10)

    def test_rectangular_B(self):
        """B is 3x1, U is 1D → BPB is 3x3."""
        B = np.array([[1], [0], [0]])
        U = Ellipsoid([0.0], np.array([[4.0]]))
        sys = LinSys(np.eye(3), B, U)
        t_eval = np.array([0.0])
        BPB, BPBsr, _, _ = precompute_bounds(sys, t_eval)
        expected = np.array([[4, 0, 0], [0, 0, 0], [0, 0, 0]], dtype=float)
        np.testing.assert_allclose(BPB[0], expected, atol=1e-10)


# =====================================================================
# precompute_all (integration test)
# =====================================================================

class TestPrecomputeAll:
    def test_smoke_lti(self):
        """Full precomputation for a simple LTI system."""
        A = np.array([[0, 1], [-1, 0]])
        B = np.array([[0], [1]])
        U = Ellipsoid([0.0], np.array([[1.0]]))
        sys = LinSys(A, B, U)
        x0 = np.array([1.0, 0.0])
        L0 = np.array([[1, 0], [0, 1]], dtype=float)
        data = precompute_all(sys, x0, L0, (0, 1), N_steps=20)

        assert data.t_eval.shape == (20,)
        assert data.Phi.shape == (20, 2, 2)
        assert data.Phi_inv.shape == (20, 2, 2)
        assert data.centers.shape == (20, 2)
        assert data.L.shape == (20, 2, 2)
        assert data.BPB.shape == (20, 2, 2)
        assert data.BPBsr.shape == (20, 2, 2)
        assert data.GQG is None

    def test_smoke_with_disturbance(self):
        """Full precomputation with disturbance channel."""
        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 = np.zeros(2)
        L0 = np.array([[1.0], [0.0]])
        data = precompute_all(sys, x0, L0, (0, 2), N_steps=10)

        assert data.GQG is not None
        assert data.GQG.shape == (10, 2, 2)
        assert data.GQGsr is not None

    def test_stm_matches_standalone(self):
        """precompute_all STM should match standalone compute_stm."""
        A = np.array([[1, 0.5], [0, -1]])
        sys = LinSys(A)
        x0 = np.array([1.0, 0.0])
        L0 = np.array([[1.0], [0.0]])
        data = precompute_all(sys, x0, L0, (0, 1), N_steps=10)
        Phi_standalone = compute_stm(sys, (0, 1), data.t_eval)
        np.testing.assert_allclose(data.Phi, Phi_standalone, atol=1e-10)


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

class TestEdgeCases:
    def test_empty_system_stm_raises(self):
        sys = LinSys()
        with pytest.raises(ValueError, match="non-empty"):
            compute_stm(sys, (0, 1), np.array([0.0]))

    def test_single_time_point(self):
        """Single time point should work."""
        sys = LinSys(np.eye(2))
        Phi = compute_stm(sys, (0, 0), np.array([0.0]))
        np.testing.assert_allclose(Phi[0], np.eye(2), atol=1e-10)

    def test_1d_system_full_pipeline(self):
        """1D system through entire pipeline."""
        A = np.array([[-2.0]])
        B = np.array([[1.0]])
        U = Ellipsoid([0.5], np.array([[0.1]]))
        sys = LinSys(A, B, U)
        x0 = np.array([1.0])
        L0 = np.array([[1.0]])
        data = precompute_all(sys, x0, L0, (0, 3), N_steps=30)
        assert data.centers.shape == (30, 1)
        # Steady state: x = -A^{-1} B p = (1/2)(1)(0.5) = 0.25
        assert data.centers[-1, 0] == pytest.approx(0.25, abs=0.05)
