"""Tests for NNTrainer (replaces test_solver.py)."""

import numpy as np
import pytest

from hpi.nn_trainer import NNTrainer


@pytest.fixture
def trainer():
    return NNTrainer(R=1.0, lr=1e-3, epochs=50, lambda_weight=1e-6)


class TestNNTrainer:
    """Test suite for NNTrainer."""

    # ── integrate_interval ──

    def test_integrate_interval_scalar(self):
        """Trapezoidal integration: [0,1,2] with dt=1 -> [0.5, 1.5]."""
        f = np.array([0.0, 1.0, 2.0])
        result = NNTrainer.integrate_interval(f, dt=1.0)
        expected = np.array([0.5, 1.5])
        np.testing.assert_allclose(result, expected, atol=1e-12)

    def test_integrate_interval_vector(self):
        """Vector trapezoidal integration."""
        f = np.array([[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]])
        result = NNTrainer.integrate_interval(f, dt=1.0)
        expected = np.array([[1.0, 2.0], [3.0, 4.0]])
        np.testing.assert_allclose(result, expected, atol=1e-12)

    # ── compute_u_hat / compute_v_hat ──

    def test_compute_u_hat_manual(self, trainer):
        """u_hat returns correct shape for single and batch input."""
        x_single = np.array([0.5, 0.3])
        u = trainer.compute_u_hat(x_single)
        assert np.isscalar(u)
        assert np.isfinite(u)

        x_batch = np.random.randn(10, 2).astype(np.float32)
        u_batch = trainer.compute_u_hat(x_batch)
        assert u_batch.shape == (10,)
        assert np.all(np.isfinite(u_batch))

    def test_compute_v_hat(self, trainer):
        """v_hat = u - u_hat."""
        u = np.array([0.5, 0.3, 0.1])
        x = np.array([[0.5, 0.3], [0.1, -0.2], [0.0, 0.0]])
        u_hat = trainer.compute_u_hat(x)
        v_hat = trainer.compute_v_hat(u, x)
        np.testing.assert_allclose(v_hat, u - u_hat, atol=1e-10)

    # ── Training output shapes ──

    def test_compute_I_phi_dim(self, trainer):
        """After training, verify that V can be computed (replaces I_phi test)."""
        M = 100
        x = np.random.randn(M, 2).astype(np.float32) * 0.5
        V = trainer.compute_V(x)
        assert V.shape == (M,)

    def test_compute_I_psi_dim(self, trainer):
        """u_hat output shape matches input (replaces I_psi test)."""
        M = 100
        x = np.random.randn(M, 2).astype(np.float32) * 0.5
        u_hat = trainer.compute_u_hat(x)
        assert u_hat.shape == (M,)

    def test_compute_I_Q_dim(self, trainer):
        """Q(x) can be computed on trajectory (replaces I_Q test)."""
        M = 100
        x = np.random.randn(M, 2).astype(np.float32) * 0.5
        Q = x[:, 0]**2 + x[:, 1]**2
        assert Q.shape == (M,)

    # ── System construction (adapted) ──

    def test_build_system_dims(self, trainer):
        """Training produces a result dict with expected keys."""
        M = 200
        t = np.linspace(0, 2.0, M)
        X = np.random.randn(M, 2).astype(np.float32) * 0.5
        U = np.random.randn(M).astype(np.float32) * 0.1

        result = trainer.train_one_iteration(t, X, U, L_i=1.0, prev_actor=None, dt=0.01)
        assert "loss_c" in result
        assert "mse_c" in result
        assert np.isfinite(result["loss_c"])

    def test_solve_dim(self, trainer):
        """Training returns result dict with finite loss."""
        M = 200
        t = np.linspace(0, 2.0, M)
        X = np.random.randn(M, 2).astype(np.float32) * 0.5
        U = np.random.randn(M).astype(np.float32) * 0.1

        result = trainer.train_one_iteration(t, X, U, L_i=1.0, prev_actor=None, dt=0.01)
        assert np.isfinite(result["loss_c"])
        assert result["loss_c"] >= 0

    def test_solve_no_nan(self, trainer):
        """Training produces no NaN in loss."""
        M = 200
        t = np.linspace(0, 2.0, M)
        X = np.random.randn(M, 2).astype(np.float32) * 0.5
        U = np.random.randn(M).astype(np.float32) * 0.1

        result = trainer.train_one_iteration(t, X, U, L_i=0.0, prev_actor=None, dt=0.01)
        assert not np.isnan(result["loss_c"])
        assert not np.isinf(result["loss_c"])

    def test_solve_with_trajectory_data(self, trainer):
        """Train on simulated pendulum data produces finite parameters."""
        from hpi.data_collector import DataCollector

        dc = DataCollector()
        x0 = np.array([0.5, 0.0])

        def u_func(t, x):
            return dc.pe_signal(t)

        t, X, U = dc.collect_trajectory(x0, T=2.0, dt=0.01, u_func=u_func)
        result = trainer.train_one_iteration(t, X, U, L_i=2.0, prev_actor=None, dt=0.01)

        assert np.isfinite(result["loss_c"])
        assert result["loss_c"] >= 0.0