"""Tests for MetaDictionary module."""

import torch
import pytest
from cymeta.dictionary import MetaDictionary, ModuleTypeDictionary


def test_meta_dictionary_initialization():
    """Test MetaDictionary initialization."""
    num_atoms = 32
    atom_dim = 128
    
    dictionary = MetaDictionary(num_atoms, atom_dim)
    
    assert dictionary.num_atoms == num_atoms
    assert dictionary.atom_dim == atom_dim
    assert dictionary.atoms.shape == (num_atoms, atom_dim)


def test_meta_dictionary_get_atom():
    """Test getting individual atoms."""
    dictionary = MetaDictionary(num_atoms=8, atom_dim=64)
    
    atom = dictionary.get_atom(0)
    assert atom.shape == (64,)
    
    # Test index bounds
    with pytest.raises(IndexError):
        dictionary.get_atom(-1)
    with pytest.raises(IndexError):
        dictionary.get_atom(8)


def test_meta_dictionary_cyclic_shift():
    """Test cyclic shift operations."""
    dictionary = MetaDictionary(num_atoms=4, atom_dim=8)
    
    # Create a test atom with known values
    with torch.no_grad():
        dictionary.atoms[0] = torch.arange(8, dtype=torch.float32)
    
    # Test forward shift
    shifted = dictionary.apply_cyclic_shift(0, shift=2)
    expected = torch.tensor([6., 7., 0., 1., 2., 3., 4., 5.])
    assert torch.allclose(shifted, expected)
    
    # Test backward shift
    shifted = dictionary.apply_cyclic_shift(0, shift=-2)
    expected = torch.tensor([2., 3., 4., 5., 6., 7., 0., 1.])
    assert torch.allclose(shifted, expected)


def test_meta_dictionary_batch_shifts():
    """Test batch cyclic shifts."""
    dictionary = MetaDictionary(num_atoms=4, atom_dim=8)
    
    atom_indices = torch.tensor([0, 1, 2])
    shifts = torch.tensor([1, -1, 2])
    
    shifted = dictionary.apply_cyclic_shifts_batch(atom_indices, shifts)
    assert shifted.shape == (3, 8)


def test_module_type_dictionary():
    """Test ModuleTypeDictionary."""
    module_dict = ModuleTypeDictionary(
        module_type="attention_q",
        num_atoms=16,
        atom_dim=256,
    )
    
    assert module_dict.module_type == "attention_q"
    assert module_dict.dictionary.num_atoms == 16
    assert module_dict.dictionary.atom_dim == 256


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

