"""Tests for DataCollector (Module A)."""

import numpy as np
import pytest

from hpi.data_collector import DataCollector


@pytest.fixture
def dc():
    return DataCollector(J=1.0, mgl=1.0)


class TestDataCollector:
    """Test suite for DataCollector."""

    def test_f_zero(self, dc):
        """f(0) = 0 at equilibrium."""
        f0 = dc.f(np.zeros(2))
        np.testing.assert_allclose(f0, [0.0, 0.0], atol=1e-12)

    def test_f_pi_half(self, dc):
        """f([π/2, 0]) = [0, mgl/J]."""
        x = np.array([np.pi / 2, 0.0])
        f_val = dc.f(x)
        expected = np.array([0.0, dc.mgl / dc.J])
        np.testing.assert_allclose(f_val, expected, atol=1e-12)

    def test_f_batch(self, dc):
        """Batch f output shape (N, 2)."""
        x = np.random.randn(10, 2)
        f_vals = dc.f(x)
        assert f_vals.shape == (10, 2)

    def test_g_constant(self, dc):
        """g(x) = [0, 1/J] constant."""
        g_val = dc.g(np.array([0.5, 0.3]))
        expected = np.array([0.0, 1.0 / dc.J])
        np.testing.assert_allclose(g_val, expected, atol=1e-12)

    def test_g_batch(self, dc):
        """Batch g output shape (N, 2)."""
        x = np.random.randn(10, 2)
        g_vals = dc.g(x)
        assert g_vals.shape == (10, 2)
        # All rows should be [0, 1/J]
        np.testing.assert_allclose(g_vals[:, 1], np.ones(10) / dc.J, atol=1e-12)

    def test_pe_signal_bounded(self, dc):
        """PE signal amplitude should be <= 0.1."""
        for t in np.linspace(0, 10, 100):
            val = dc.pe_signal(t)
            assert abs(val) <= 0.1 + 1e-12

    def test_pe_signal_nonzero(self, dc):
        """PE signal should not be identically zero."""
        vals = [dc.pe_signal(t) for t in np.linspace(0, 10, 100)]
        assert np.max(np.abs(vals)) > 1e-6

    def test_dynamics_manual(self, dc):
        """Verify ẋ = f(x) + g(x)·u."""
        x = np.array([0.5, 0.3])
        u = 0.7
        dx = dc.dynamics(0.0, x, lambda t, x: u)
        expected = dc.f(x) + dc.g(x) * u
        np.testing.assert_allclose(dx, expected, atol=1e-12)

    def test_collect_trajectory_shape(self, dc):
        """Trajectory outputs have consistent dimensions."""
        x0 = np.array([0.5, 0.0])
        T = 1.0
        dt = 0.01

        def u_const(t, x):
            return 0.0

        t, X, U = dc.collect_trajectory(x0, T, dt, u_const)

        M_expected = int(T / dt)
        assert t.shape == (M_expected,)
        assert X.shape == (M_expected, 2)
        assert U.shape == (M_expected,)

    def test_collect_trajectory_start(self, dc):
        """First state should equal x0."""
        x0 = np.array([0.5, 0.0])
        T = 1.0
        dt = 0.01

        def u_const(t, x):
            return 0.0

        t, X, U = dc.collect_trajectory(x0, T, dt, u_const)
        np.testing.assert_allclose(X[0], x0, atol=1e-10)

    def test_collect_trajectory_drift(self, dc):
        """With zero control, state should evolve (pendulum swings)."""
        x0 = np.array([0.5, 0.0])
        T = 2.0
        dt = 0.01

        def u_const(t, x):
            return 0.0

        t, X, U = dc.collect_trajectory(x0, T, dt, u_const)
        # Should not stay at x0 (pendulum moves)
        assert np.linalg.norm(X[-1] - x0) > 1e-6