"""Tests for GatingNetwork module."""

import torch
import pytest
from cymeta.gating import GatingNetwork, ContextualGating


def test_gating_network_initialization():
    """Test GatingNetwork initialization."""
    gating = GatingNetwork(
        input_dim=128,
        num_atoms=32,
        hidden_dim=16,
    )
    
    assert gating.input_dim == 128
    assert gating.num_atoms == 32
    assert gating.hidden_dim == 16


def test_gating_network_forward():
    """Test GatingNetwork forward pass."""
    gating = GatingNetwork(
        input_dim=64,
        num_atoms=16,
        hidden_dim=8,
    )
    
    # Single token
    token_embedding = torch.randn(4, 64)  # (batch_size, input_dim)
    gating_coeffs = gating(token_embedding)
    
    assert gating_coeffs.shape == (4, 16)
    # Check that coefficients sum to 1 (softmax normalization)
    assert torch.allclose(gating_coeffs.sum(dim=-1), torch.ones(4), atol=1e-5)
    
    # Sequence
    token_embedding_seq = torch.randn(2, 10, 64)  # (batch_size, seq_len, input_dim)
    gating_coeffs_seq = gating(token_embedding_seq)
    
    assert gating_coeffs_seq.shape == (2, 16)


def test_gating_network_head_specific():
    """Test head-specific gating network."""
    gating = GatingNetwork(
        input_dim=128,
        num_atoms=32,
        hidden_dim=16,
        num_heads=8,
        use_head_specific=True,
    )
    
    token_embedding = torch.randn(4, 128)
    
    # Test with head index
    gating_coeffs = gating(token_embedding, head=0)
    assert gating_coeffs.shape == (4, 32)
    
    # Test invalid head
    with pytest.raises(ValueError):
        gating(token_embedding, head=8)


def test_gating_network_num_parameters():
    """Test parameter counting."""
    gating = GatingNetwork(
        input_dim=64,
        num_atoms=16,
        hidden_dim=8,
    )
    
    num_params = gating.get_num_parameters()
    assert num_params > 0
    # Should be small (lightweight network)
    assert num_params < 10000


def test_contextual_gating():
    """Test ContextualGating."""
    gating = ContextualGating(
        input_dim=128,
        num_atoms=32,
        hidden_dim=16,
        context_dim=64,
    )
    
    token_embedding = torch.randn(4, 128)
    context = torch.randn(4, 64)
    
    gating_coeffs = gating(token_embedding, context)
    assert gating_coeffs.shape == (4, 32)


if __name__ == "__main__":
    pytest.main([__file__])

