"""Tests for the Neon VLA model (CPU-only, no backbone loading)."""

import numpy as np
import pytest
import torch

from neon.data.action_space import ControlMode, G1ActionSpace
from neon.model.action_heads import (
    ActionChunkingHead, ActionHeadConfig, G1ActionHead, MLPDecoder,
    ReLUSquared, RMSNorm, soft_cap,
)
from neon.model.neon_vla import NeonConfig, NeonVLA, NeonOutput


# =============================================================================
# Parameter Golf v2 primitives
# =============================================================================


class TestReLUSquared:
    """Test ReLU² activation."""

    def test_positive_values(self):
        """Positive inputs → squared."""
        act = ReLUSquared()
        x = torch.tensor([1.0, 2.0, 3.0])
        out = act(x)
        expected = torch.tensor([1.0, 4.0, 9.0])
        torch.testing.assert_close(out, expected)

    def test_negative_values(self):
        """Negative inputs → zero."""
        act = ReLUSquared()
        x = torch.tensor([-1.0, -2.0, -0.5])
        out = act(x)
        assert (out == 0).all()

    def test_gradient_continuity(self):
        """Gradient at x=0 should be 0 (continuous, unlike ReLU)."""
        act = ReLUSquared()
        x = torch.tensor([0.0], requires_grad=True)
        y = act(x)
        y.backward()
        assert x.grad.item() == 0.0  # Continuous at zero


class TestRMSNorm:
    """Test RMSNorm layer."""

    def test_output_shape(self):
        """Should preserve shape."""
        norm = RMSNorm(64)
        x = torch.randn(4, 64)
        out = norm(x)
        assert out.shape == (4, 64)

    def test_normalization(self):
        """RMS of output should be approximately 1."""
        norm = RMSNorm(256)
        x = torch.randn(8, 256) * 10  # Large values
        out = norm(x)
        rms = out.pow(2).mean(-1).sqrt()
        # RMS should be close to 1 (weight is initialized to 1)
        torch.testing.assert_close(rms, torch.ones_like(rms), atol=0.1, rtol=0.1)

    def test_zero_input(self):
        """Should handle zero input without NaN."""
        norm = RMSNorm(32)
        x = torch.zeros(2, 32)
        out = norm(x)
        assert not out.isnan().any()


class TestSoftCap:
    """Test logit soft-capping."""

    def test_small_values_passthrough(self):
        """Small values should be approximately unchanged."""
        x = torch.tensor([0.1, -0.1, 0.5])
        out = soft_cap(x, cap=1.0)
        # For small x relative to cap, tanh(x) ≈ x
        torch.testing.assert_close(out, x, atol=0.05, rtol=0.1)

    def test_large_values_clamped(self):
        """Large values should approach ±cap."""
        x = torch.tensor([100.0, -100.0])
        out = soft_cap(x, cap=1.0)
        assert out[0].item() <= 1.0 and out[0].item() > 0.99
        assert out[1].item() >= -1.0 and out[1].item() < -0.99

    def test_smoother_than_tanh(self):
        """Gradient near boundary should be larger than hard tanh."""
        x = torch.tensor([0.95], requires_grad=True)
        y = soft_cap(x, cap=1.0)
        y.backward()
        grad_soft = x.grad.item()

        x2 = torch.tensor([0.95], requires_grad=True)
        y2 = torch.tanh(x2)
        y2.backward()
        grad_tanh = x2.grad.item()

        # soft_cap has gentler saturation than raw tanh
        # Both should have positive gradient
        assert grad_soft > 0
        assert grad_tanh > 0


class TestMLPDecoder:
    """Test the MLP decoder with v2 enhancements."""

    def test_output_shape(self):
        """MLP should produce correct output shape."""
        decoder = MLPDecoder(input_dim=256, output_dim=17)
        x = torch.randn(4, 256)
        out = decoder(x)
        assert out.shape == (4, 17)

    def test_output_range(self):
        """MLP with soft-cap should output in approximately [-1, 1]."""
        decoder = MLPDecoder(input_dim=128, output_dim=10)
        x = torch.randn(8, 128) * 10  # Large inputs
        out = decoder(x)
        assert out.min() >= -1.0
        assert out.max() <= 1.0

    def test_legacy_tanh_mode(self):
        """soft_cap_value=0 should use hard Tanh."""
        decoder = MLPDecoder(input_dim=64, output_dim=5, soft_cap_value=0.0)
        x = torch.randn(2, 64) * 100
        out = decoder(x)
        assert out.min() >= -1.0
        assert out.max() <= 1.0

    def test_v2_disabled(self):
        """All v2 features can be disabled for backward compat."""
        decoder = MLPDecoder(
            input_dim=64, output_dim=5, num_layers=3,
            use_relu_squared=False, use_rmsnorm=False,
            use_residual_scale=False, use_skip_connections=False,
            soft_cap_value=0.0,
        )
        x = torch.randn(2, 64)
        out = decoder(x)
        assert out.shape == (2, 5)

    def test_skip_connections_require_depth(self):
        """Skip connections only activate with num_layers >= 3."""
        # 2-layer MLP — skip should be disabled internally
        decoder = MLPDecoder(
            input_dim=64, output_dim=5, num_layers=2,
            use_skip_connections=True,
        )
        x = torch.randn(2, 64)
        out = decoder(x)
        assert out.shape == (2, 5)

    def test_residual_scales_learnable(self):
        """Residual scales should be parameters that can be optimized."""
        decoder = MLPDecoder(
            input_dim=64, output_dim=5, num_layers=4,
            use_residual_scale=True,
        )
        assert decoder.residual_scales is not None
        # All scales start at 1.0
        for s in decoder.residual_scales:
            assert s.item() == 1.0
            assert s.requires_grad


class TestActionChunkingHead:
    """Test the action chunking decoder."""

    def test_output_shape(self):
        """Chunking head should produce (batch, steps, dim)."""
        head = ActionChunkingHead(input_dim=256, action_dim=17, num_steps=16)
        x = torch.randn(4, 256)
        out = head(x)
        assert out.shape == (4, 16, 17)

    def test_output_range(self):
        """Chunking head with Tanh should output in [-1, 1]."""
        head = ActionChunkingHead(input_dim=128, action_dim=10, num_steps=8)
        x = torch.randn(2, 128)
        out = head(x)
        assert out.min() >= -1.0
        assert out.max() <= 1.0


class TestG1ActionHead:
    """Test the G1-specific multi-head decoder."""

    def test_single_head(self):
        """Single head mode should work."""
        space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=True)
        config = ActionHeadConfig(
            hidden_size=256,
            action_dim=space.action_dim,
            use_separate_heads=False,
            use_action_chunking=False,
        )
        head = G1ActionHead(config, space)
        x = torch.randn(4, 256)
        out = head(x)
        assert out.shape == (4, space.action_dim)

    def test_separate_heads_arms_only(self):
        """Separate heads in arms_only mode."""
        space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=True)
        config = ActionHeadConfig(
            hidden_size=256,
            action_dim=space.action_dim,
            use_separate_heads=True,
            use_action_chunking=False,
        )
        head = G1ActionHead(config, space)
        x = torch.randn(4, 256)
        out = head(x)
        assert out.shape[0] == 4
        assert out.shape[1] == space.action_dim

    def test_chunking_with_separate_heads(self):
        """Separate heads with action chunking."""
        space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=True)
        config = ActionHeadConfig(
            hidden_size=256,
            action_dim=space.action_dim,
            num_action_steps=8,
            use_separate_heads=True,
            use_action_chunking=True,
        )
        head = G1ActionHead(config, space)
        x = torch.randn(2, 256)
        out = head(x)
        assert out.shape == (2, 8, space.action_dim)

    def test_loss_computation(self):
        """Loss should be a scalar."""
        space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=False)
        config = ActionHeadConfig(
            hidden_size=128,
            action_dim=space.action_dim,
            use_separate_heads=False,
            use_action_chunking=False,
        )
        head = G1ActionHead(config, space)
        predicted = torch.randn(4, space.action_dim)
        target = torch.randn(4, space.action_dim)
        loss = head.compute_loss(predicted, target)
        assert loss.ndim == 0  # Scalar


class TestNeonVLA:
    """Test the full NeonVLA model (without backbone loading)."""

    def test_config_serialization(self):
        """Config should serialize and deserialize."""
        import tempfile
        import os

        config = NeonConfig(control_mode="arms_only", include_locomotion=True)
        path = os.path.join(tempfile.gettempdir(), "test_neon_config.json")
        config.save(path)
        loaded = NeonConfig.load(path)
        assert loaded.control_mode == "arms_only"
        assert loaded.include_locomotion is True
        os.unlink(path)

    def test_model_creation(self):
        """Model should be creatable without loading backbone."""
        config = NeonConfig(
            control_mode="arms_only",
            include_locomotion=True,
            num_action_steps=8,
        )
        model = NeonVLA(config)
        assert model.action_space.mode == ControlMode.ARMS_ONLY
        assert model.action_space.action_dim == 17  # 14 arms + 3 locomotion

    def test_action_head_forward(self):
        """Action heads should work with mock features."""
        config = NeonConfig(
            control_mode="arms_only",
            include_locomotion=False,
            num_action_steps=4,
            use_proprioception=False,  # Disable proprio so fusion input == hidden_size
        )
        model = NeonVLA(config)

        # Mock backbone output (must match fusion layer input)
        mock_features = torch.randn(2, model.backbone.hidden_size)

        # Test fusion + action head
        fused = model.fusion(mock_features)
        actions = model.action_head(fused)
        assert actions.shape == (2, 4, 14)  # batch=2, steps=4, 14 arm joints

    def test_neon_output(self):
        """NeonOutput should store results correctly."""
        output = NeonOutput(
            actions=np.random.randn(8, 17).astype(np.float32),
            raw_actions=np.random.randn(8, 17).astype(np.float32),
        )
        assert output.actions.shape == (8, 17)
