"""Unit tests for the LinSys class (Phase 3A).

Covers: construction, validation, time-varying matrices, bounds (U/V/W),
query methods (is_lti, has_disturbance, has_noise, dimension), display,
and evaluation at time t.
"""

import numpy as np
import pytest

from ellipsoidal_toolbox.core.ellipsoid import Ellipsoid
from ellipsoidal_toolbox.systems.linsys import LinSys


# =====================================================================
# Construction — empty
# =====================================================================

class TestEmptyConstruction:
    def test_no_args(self):
        S = LinSys()
        assert S.is_empty()
        assert S.dimension() == (0, 0, 0, 0)
        assert S.is_lti
        assert not S.is_discrete

    def test_repr_empty(self):
        assert "empty" in repr(LinSys())

    def test_str_empty(self):
        assert "Empty" in str(LinSys())


# =====================================================================
# Construction — LTI systems
# =====================================================================

class TestLTIConstruction:
    def test_A_only(self):
        """Autonomous system dx/dt = Ax."""
        S = LinSys(np.eye(3))
        assert not S.is_empty()
        assert S.is_lti
        n, k, m, d = S.dimension()
        assert n == 3
        assert k == 0
        assert m == 3  # default C = I
        assert d == 0

    def test_A_and_B(self):
        """System with control input, no bounds."""
        A = np.array([[0, 1], [-1, 0]])
        B = np.array([[0], [1]])
        S = LinSys(A, B)
        n, k, m, d = S.dimension()
        assert n == 2
        assert k == 1
        assert m == 2
        assert d == 0
        assert S.is_lti

    def test_full_system(self):
        """dx/dt = Ax + Bu + Gv, y = Cx + w, with all bounds."""
        A = np.eye(2)
        B = np.ones((2, 1))
        U = Ellipsoid([0.0], np.eye(1))
        G = np.eye(2)
        V = Ellipsoid([0, 0], 0.1 * np.eye(2))
        C = np.array([[1, 0]])
        W = Ellipsoid([0.0], 0.01 * np.eye(1))
        S = LinSys(A, B, U, G, V, C, W)
        n, k, m, d = S.dimension()
        assert n == 2
        assert k == 1
        assert m == 1
        assert d == 2
        assert S.is_lti
        assert S.has_disturbance()
        assert S.has_noise()

    def test_discrete(self):
        S = LinSys(np.eye(2), discrete=True)
        assert S.is_discrete

    def test_continuous_default(self):
        S = LinSys(np.eye(2))
        assert not S.is_discrete

    def test_no_disturbance(self):
        S = LinSys(np.eye(2), np.ones((2, 1)))
        assert not S.has_disturbance()

    def test_no_noise(self):
        S = LinSys(np.eye(2))
        assert not S.has_noise()


# =====================================================================
# Construction — time-varying (LTV) systems
# =====================================================================

class TestLTVConstruction:
    def test_callable_A(self):
        """Time-varying A makes system LTV."""
        def A_func(t):
            return np.array([[0, 1], [-1 - t, 0]])

        S = LinSys(A_func)
        assert not S.is_lti
        n, _, _, _ = S.dimension()
        assert n == 2

    def test_callable_B(self):
        def B_func(t):
            return np.array([[0], [1 + t]])

        S = LinSys(np.eye(2), B_func)
        assert not S.is_lti

    def test_callable_G(self):
        def G_func(t):
            return np.eye(2) * np.exp(-t)

        S = LinSys(np.eye(2), G=G_func)
        assert not S.is_lti

    def test_callable_C(self):
        def C_func(t):
            return np.array([[1, t]])

        S = LinSys(np.eye(2), C=C_func)
        assert not S.is_lti

    def test_time_varying_bound_U(self):
        """Callable U is time-varying."""
        def U_func(t):
            return Ellipsoid([0.0], np.eye(1) * (1 + t))

        S = LinSys(np.eye(2), np.ones((2, 1)), U=U_func)
        assert not S.is_empty()

    def test_time_varying_bound_V(self):
        def V_func(t):
            return Ellipsoid([0, 0], np.eye(2) * (1 + t))

        S = LinSys(np.eye(2), G=np.eye(2), V=V_func)
        assert S.has_disturbance()


# =====================================================================
# Evaluation at time t
# =====================================================================

class TestEvaluation:
    def test_A_at_constant(self):
        A = np.array([[1, 2], [3, 4]])
        S = LinSys(A)
        np.testing.assert_array_equal(S.A_at(0.0), A)
        np.testing.assert_array_equal(S.A_at(5.0), A)

    def test_A_at_callable(self):
        def A_func(t):
            return np.array([[t, 0], [0, -t]])

        S = LinSys(A_func)
        np.testing.assert_array_equal(S.A_at(2.0), [[2, 0], [0, -2]])

    def test_B_at_none(self):
        S = LinSys(np.eye(2))
        assert S.B_at(0.0) is None

    def test_B_at_constant(self):
        B = np.array([[1], [0]])
        S = LinSys(np.eye(2), B)
        np.testing.assert_array_equal(S.B_at(3.0), B)

    def test_U_at_ellipsoid(self):
        U = Ellipsoid([0.0], np.eye(1))
        S = LinSys(np.eye(2), np.ones((2, 1)), U=U)
        U_eval = S.U_at(0.0)
        assert isinstance(U_eval, Ellipsoid)
        assert not U_eval.is_empty()

    def test_U_at_none(self):
        S = LinSys(np.eye(2), np.ones((2, 1)))
        assert S.U_at(0.0) is None

    def test_U_at_fixed_vector(self):
        """Fixed vector bound creates a zero-shape ellipsoid."""
        S = LinSys(np.eye(2), np.ones((2, 1)), U=np.array([1.0]))
        U_eval = S.U_at(0.0)
        assert isinstance(U_eval, Ellipsoid)
        np.testing.assert_array_almost_equal(U_eval.center, [1.0])

    def test_U_at_callable(self):
        def U_func(t):
            return Ellipsoid([t], np.eye(1) * (1 + t))

        S = LinSys(np.eye(2), np.ones((2, 1)), U=U_func)
        U2 = S.U_at(2.0)
        assert isinstance(U2, Ellipsoid)
        np.testing.assert_array_almost_equal(U2.center, [2.0])

    def test_V_at_with_disturbance(self):
        V = Ellipsoid([0, 0], 0.5 * np.eye(2))
        S = LinSys(np.eye(2), G=np.eye(2), V=V)
        V_eval = S.V_at(0.0)
        assert isinstance(V_eval, Ellipsoid)

    def test_W_at_with_noise(self):
        W = Ellipsoid([0, 0], 0.01 * np.eye(2))
        S = LinSys(np.eye(2), W=W)
        W_eval = S.W_at(0.0)
        assert isinstance(W_eval, Ellipsoid)

    def test_G_at_constant(self):
        G = np.eye(3)
        S = LinSys(np.eye(3), G=G)
        np.testing.assert_array_equal(S.G_at(1.0), G)

    def test_C_at_default_identity(self):
        S = LinSys(np.eye(3))
        np.testing.assert_array_equal(S.C_at(0.0), np.eye(3))


# =====================================================================
# Validation
# =====================================================================

class TestValidation:
    def test_A_not_square_raises(self):
        with pytest.raises(ValueError, match="square"):
            LinSys(np.ones((2, 3)))

    def test_B_row_mismatch_raises(self):
        with pytest.raises(ValueError, match="row count"):
            LinSys(np.eye(2), np.ones((3, 1)))

    def test_G_row_mismatch_raises(self):
        with pytest.raises(ValueError, match="row count"):
            LinSys(np.eye(2), G=np.ones((3, 1)))

    def test_C_col_mismatch_raises(self):
        with pytest.raises(ValueError, match="column count"):
            LinSys(np.eye(2), C=np.ones((1, 3)))

    def test_U_dim_mismatch_raises(self):
        with pytest.raises(ValueError, match="dimension"):
            LinSys(np.eye(2), np.ones((2, 1)), U=Ellipsoid([0, 0], np.eye(2)))

    def test_V_dim_mismatch_raises(self):
        with pytest.raises(ValueError, match="dimension"):
            LinSys(np.eye(2), G=np.eye(2), V=Ellipsoid([0], np.eye(1)))

    def test_W_dim_mismatch_raises(self):
        with pytest.raises(ValueError, match="dimension"):
            LinSys(np.eye(2), W=Ellipsoid([0], np.eye(1)))

    def test_U_vector_dim_mismatch_raises(self):
        with pytest.raises(ValueError, match="vector length"):
            LinSys(np.eye(2), np.ones((2, 1)), U=np.array([1.0, 2.0]))

    def test_callable_B_row_mismatch_raises(self):
        def B_bad(t):
            return np.ones((5, 1))

        with pytest.raises(ValueError, match="row count"):
            LinSys(np.eye(2), B_bad)

    def test_callable_G_row_mismatch_raises(self):
        def G_bad(t):
            return np.ones((5, 1))

        with pytest.raises(ValueError, match="row count"):
            LinSys(np.eye(2), G=G_bad)

    def test_callable_C_col_mismatch_raises(self):
        def C_bad(t):
            return np.ones((1, 5))

        with pytest.raises(ValueError, match="column count"):
            LinSys(np.eye(2), C=C_bad)


# =====================================================================
# Properties and queries
# =====================================================================

class TestProperties:
    def test_A_property(self):
        A = np.eye(2)
        S = LinSys(A)
        np.testing.assert_array_equal(S.A, A)

    def test_B_property_none(self):
        S = LinSys(np.eye(2))
        assert S.B is None

    def test_control_property(self):
        U = Ellipsoid([0.0], np.eye(1))
        S = LinSys(np.eye(2), np.ones((2, 1)), U=U)
        assert isinstance(S.control, Ellipsoid)

    def test_disturbance_property(self):
        V = Ellipsoid([0, 0], np.eye(2))
        S = LinSys(np.eye(2), G=np.eye(2), V=V)
        assert isinstance(S.disturbance, Ellipsoid)

    def test_noise_property(self):
        W = Ellipsoid([0, 0], np.eye(2))
        S = LinSys(np.eye(2), W=W)
        assert isinstance(S.noise, Ellipsoid)


# =====================================================================
# Display
# =====================================================================

class TestDisplay:
    def test_repr_lti(self):
        S = LinSys(np.eye(2), np.ones((2, 1)))
        r = repr(S)
        assert "continuous" in r
        assert "LTI" in r
        assert "state=2" in r

    def test_repr_ltv(self):
        S = LinSys(lambda t: np.eye(2) * t)
        assert "LTV" in repr(S)

    def test_repr_discrete(self):
        S = LinSys(np.eye(2), discrete=True)
        assert "discrete" in repr(S)

    def test_str_full(self):
        A = np.eye(2)
        B = np.ones((2, 1))
        U = Ellipsoid([0.0], np.eye(1))
        G = np.eye(2)
        V = Ellipsoid([0, 0], 0.1 * np.eye(2))
        S = LinSys(A, B, U, G, V)
        s = str(S)
        assert "state dim" in s
        assert "input dim" in s
        assert "disturbance dim" in s
        assert "control bounds" in s

    def test_str_empty(self):
        assert "Empty" in str(LinSys())


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

class TestEdgeCases:
    def test_G_without_V(self):
        """G matrix without V bounds = unbounded disturbance channel."""
        S = LinSys(np.eye(2), G=np.eye(2))
        _, _, _, d = S.dimension()
        assert d == 2
        assert not S.has_disturbance()  # V is None

    def test_empty_ellipsoid_bound(self):
        """Passing empty Ellipsoid as bound treated as None."""
        S = LinSys(np.eye(2), np.ones((2, 1)), U=Ellipsoid())
        assert S.control is None

    def test_1d_system(self):
        S = LinSys(np.array([[2.0]]), np.array([[1.0]]),
                    U=Ellipsoid([0.0], np.array([[1.0]])))
        n, k, m, d = S.dimension()
        assert n == 1
        assert k == 1

    def test_high_dim_system(self):
        n = 10
        S = LinSys(np.eye(n), np.ones((n, 3)),
                    U=Ellipsoid(np.zeros(3), np.eye(3)))
        dims = S.dimension()
        assert dims == (10, 3, 10, 0)

    def test_B_as_1d_array(self):
        """A 1D array B is reshaped to column vector."""
        S = LinSys(np.eye(2), np.array([1.0, 0.0]))
        B = S.B_at(0.0)
        assert B.shape == (2, 1)

    def test_C_as_1d_array(self):
        """A 1D array C is reshaped to row vector."""
        S = LinSys(np.eye(2), C=np.array([1.0, 0.0]))
        C = S.C_at(0.0)
        assert C.shape == (1, 2)
