"""Tests for NN models: CriticNN, ActorNN, compute_q, check_positive_definite."""

import numpy as np
import pytest
import torch

from hpi.nn_models import ActorNN, CriticNN, check_positive_definite, compute_q


@pytest.fixture
def critic():
    return CriticNN()


@pytest.fixture
def actor():
    return ActorNN()


class TestCriticNN:
    """Test suite for CriticNN."""

    def test_phi_zero(self, critic):
        """V(0) = 0."""
        x = torch.zeros(1, 2)
        V = critic(x)
        assert abs(V.item()) < 1e-6

    def test_phi_dim_single(self, critic):
        """Single sample: V output shape (,)."""
        x = torch.tensor([[0.5, -0.3]], dtype=torch.float32)
        V = critic(x)
        assert V.shape == (1,)
        assert V.ndim == 1

    def test_phi_dim_batch(self, critic):
        """Batch input: V output shape (N,)."""
        x = torch.randn(10, 2)
        V = critic(x)
        assert V.shape == (10,)

    def test_phi_nonnegative(self, critic):
        """V(x) >= 0 for all x."""
        x = torch.randn(100, 2)
        V = critic(x)
        assert (V >= -1e-6).all()


class TestActorNN:
    """Test suite for ActorNN."""

    def test_psi_zero(self, actor):
        """u(0) is finite (not NaN, not inf)."""
        x = torch.zeros(1, 2)
        u = actor(x)
        assert torch.isfinite(u).all()

    def test_psi_dim_single(self, actor):
        """Single sample: u output shape (,)."""
        x = torch.tensor([[0.5, -0.3]], dtype=torch.float32)
        u = actor(x)
        assert u.shape == (1,)
        assert u.ndim == 1

    def test_psi_dim_batch(self, actor):
        """Batch input: u output shape (N,)."""
        x = torch.randn(10, 2)
        u = actor(x)
        assert u.shape == (10,)

    def test_psi_exact_values(self, actor):
        """u varies with input (not constant)."""
        x1 = torch.tensor([[0.5, 0.3]], dtype=torch.float32)
        x2 = torch.tensor([[-0.8, 1.2]], dtype=torch.float32)
        u1 = actor(x1).item()
        u2 = actor(x2).item()
        assert abs(u1 - u2) > 1e-8


class TestCriticGrad:
    """Test suite for CriticNN gradient (replaces evaluate_dphi tests)."""

    @pytest.fixture
    def critic(self):
        return CriticNN()

    def test_dphi_dim_single(self, critic):
        """Single sample: grad_V shape (2,)."""
        x = torch.tensor([[0.5, -0.3]], dtype=torch.float32, requires_grad=True)
        V = critic(x)
        grad_V = torch.autograd.grad(V.sum(), x)[0]
        assert grad_V.shape == (1, 2)

    def test_dphi_dim_batch(self, critic):
        """Batch: grad_V shape (N, 2)."""
        x = torch.randn(10, 2, requires_grad=True)
        V = critic(x)
        grad_V = torch.autograd.grad(V.sum(), x)[0]
        assert grad_V.shape == (10, 2)

    def test_dphi_finite_difference(self, critic):
        """grad_V matches finite difference approximation."""
        x_np = np.array([[0.5, 0.3]], dtype=np.float32)
        x = torch.tensor(x_np, requires_grad=True)
        V = critic(x)
        grad_V = torch.autograd.grad(V.sum(), x)[0].detach().numpy()[0]

        eps = 1e-4
        for i in range(2):
            dx = np.zeros(2)
            dx[i] = eps
            x_plus = torch.tensor(x_np + dx.reshape(1, 2), dtype=torch.float32)
            x_minus = torch.tensor(x_np - dx.reshape(1, 2), dtype=torch.float32)
            with torch.no_grad():
                V_plus = critic(x_plus).item()
                V_minus = critic(x_minus).item()
            dphi_fd = (V_plus - V_minus) / (2 * eps)
            np.testing.assert_allclose(grad_V[i], dphi_fd, rtol=1e-3, atol=1e-4)

    def test_dphi_finite_difference_batch(self, critic):
        """Batch finite difference check for grad_V."""
        np.random.seed(123)
        x_np = np.random.randn(5, 2).astype(np.float32)
        x = torch.tensor(x_np, requires_grad=True)
        V = critic(x)
        grad_V = torch.autograd.grad(V.sum(), x)[0].detach().numpy()

        eps = 1e-4
        for j in range(5):
            for i in range(2):
                dx = np.zeros(2)
                dx[i] = eps
                x_plus = x_np.copy()
                x_minus = x_np.copy()
                x_plus[j] += dx
                x_minus[j] -= dx
                with torch.no_grad():
                    V_plus = critic(torch.tensor(x_plus)).numpy()[j]
                    V_minus = critic(torch.tensor(x_minus)).numpy()[j]
                dphi_fd = (V_plus - V_minus) / (2 * eps)
                np.testing.assert_allclose(grad_V[j, i], dphi_fd, rtol=1e-2, atol=1e-3)


class TestComputeQ:
    """Test suite for compute_q."""

    def test_Q_value(self):
        """Q(x) = x1^2 + x2^2."""
        x = torch.tensor([[0.6, -0.8]], dtype=torch.float32)
        q = compute_q(x)
        expected = 0.6**2 + (-0.8)**2
        assert abs(q.item() - expected) < 1e-6

    def test_Q_batch_dim(self):
        """Batch Q output shape (N,)."""
        x = torch.randn(10, 2)
        q = compute_q(x)
        assert q.shape == (10,)

    def test_Q_batch_values(self):
        """Batch Q values match manual computation."""
        x = torch.tensor([[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]], dtype=torch.float32)
        q = compute_q(x)
        expected = torch.tensor([5.0, 25.0, 0.0])
        assert torch.allclose(q, expected)