"""Tests for IndexMap module."""

import torch
import pytest
import tempfile
import os
from cymeta.index_map import IndexMap


def test_index_map_initialization():
    """Test IndexMap initialization."""
    out_dim = 64
    in_dim = 128
    num_atoms = 32
    
    atom_indices = torch.randint(0, num_atoms, (out_dim, in_dim), dtype=torch.int32)
    shift_indices = torch.randint(-64, 65, (out_dim, in_dim), dtype=torch.int32)
    
    index_map = IndexMap(atom_indices, shift_indices, max_shift=64)
    
    assert index_map.shape == (out_dim, in_dim)
    assert index_map.max_shift == 64


def test_index_map_create_empty():
    """Test creating empty index map."""
    index_map = IndexMap.create_empty(
        out_dim=32,
        in_dim=64,
        num_atoms=16,
        max_shift=32,
    )
    
    assert index_map.shape == (32, 64)
    assert torch.all(index_map.atom_indices >= 0)
    assert torch.all(index_map.atom_indices < 16)
    assert torch.all(torch.abs(index_map.shift_indices) <= 32)


def test_index_map_get_indices():
    """Test getting indices for a row."""
    atom_indices = torch.randint(0, 8, (10, 20), dtype=torch.int32)
    shift_indices = torch.randint(-10, 11, (10, 20), dtype=torch.int32)
    index_map = IndexMap(atom_indices, shift_indices)
    
    atom_row, shift_row = index_map.get_indices(0)
    assert atom_row.shape == (20,)
    assert shift_row.shape == (20,)
    assert torch.equal(atom_row, atom_indices[0])
    assert torch.equal(shift_row, shift_indices[0])


def test_index_map_get_atom_index():
    """Test getting atom index for specific position."""
    atom_indices = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32)
    shift_indices = torch.zeros(2, 3, dtype=torch.int32)
    index_map = IndexMap(atom_indices, shift_indices)
    
    assert index_map.get_atom_index(0, 0) == 0
    assert index_map.get_atom_index(0, 1) == 1
    assert index_map.get_atom_index(1, 2) == 5


def test_index_map_save_load():
    """Test saving and loading index map."""
    atom_indices = torch.randint(0, 16, (32, 64), dtype=torch.int32)
    shift_indices = torch.randint(-32, 33, (32, 64), dtype=torch.int32)
    index_map = IndexMap(atom_indices, shift_indices, max_shift=32)
    
    with tempfile.TemporaryDirectory() as tmpdir:
        filepath = os.path.join(tmpdir, "test_index_map.pkl")
        
        # Save
        index_map.save(filepath, compress=False)
        
        # Load
        loaded = IndexMap.load(filepath)
        
        assert loaded.shape == index_map.shape
        assert loaded.max_shift == index_map.max_shift
        assert torch.equal(loaded.atom_indices, index_map.atom_indices)
        assert torch.equal(loaded.shift_indices, index_map.shift_indices)


def test_index_map_compression_ratio():
    """Test compression ratio calculation."""
    index_map = IndexMap.create_empty(64, 128, num_atoms=32)
    ratio = index_map.get_compression_ratio()
    
    # Index maps alone don't compress (they're part of larger scheme)
    # They use int32 (4 bytes) same as float32, so ratio is typically < 1.0
    # This is expected - compression comes from dictionary + gating network
    assert ratio > 0.0  # Should return a valid ratio


def test_index_map_device():
    """Test moving index map to different devices."""
    index_map = IndexMap.create_empty(32, 64, num_atoms=16)
    
    # Move to CPU (should work even if already on CPU)
    cpu_map = index_map.cpu()
    assert cpu_map.atom_indices.device.type == "cpu"
    
    # Move to CUDA if available
    if torch.cuda.is_available():
        cuda_map = index_map.cuda()
        assert cuda_map.atom_indices.device.type == "cuda"


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

