"""Tests for weight reconstruction module."""

import torch
import pytest
from cymeta.dictionary import MetaDictionary
from cymeta.index_map import IndexMap
from cymeta.gating import GatingNetwork
from cymeta.reconstruct import (
    reconstruct_weights,
    reconstruct_weights_vectorized,
    reconstruct_row,
)


def test_reconstruct_weights():
    """Test weight reconstruction."""
    in_dim = 64
    out_dim = 32
    num_atoms = 16
    
    dictionary = MetaDictionary(num_atoms, in_dim)
    index_map = IndexMap.create_empty(out_dim, in_dim, num_atoms)
    gating_network = GatingNetwork(
        input_dim=in_dim,
        num_atoms=num_atoms,
        hidden_dim=8,
    )
    
    token_embedding = torch.randn(1, in_dim)
    
    reconstructed = reconstruct_weights(
        dictionary, index_map, gating_network, token_embedding
    )
    
    assert reconstructed.shape == (out_dim, in_dim)


def test_reconstruct_row():
    """Test single row reconstruction."""
    in_dim = 32
    num_atoms = 8
    
    dictionary = MetaDictionary(num_atoms, in_dim)
    atom_indices = torch.randint(0, num_atoms, (in_dim,), dtype=torch.int32)
    shift_indices = torch.randint(-16, 17, (in_dim,), dtype=torch.int32)
    gating_coeffs = torch.softmax(torch.randn(num_atoms), dim=0)
    
    row = reconstruct_row(dictionary, atom_indices, shift_indices, gating_coeffs)
    
    assert row.shape == (in_dim,)


def test_reconstruct_weights_vectorized():
    """Test vectorized weight reconstruction."""
    in_dim = 64
    out_dim = 32
    num_atoms = 16
    
    dictionary = MetaDictionary(num_atoms, in_dim)
    index_map = IndexMap.create_empty(out_dim, in_dim, num_atoms)
    gating_network = GatingNetwork(
        input_dim=in_dim,
        num_atoms=num_atoms,
        hidden_dim=8,
    )
    
    token_embedding = torch.randn(1, in_dim)
    
    reconstructed = reconstruct_weights_vectorized(
        dictionary, index_map, gating_network, token_embedding
    )
    
    assert reconstructed.shape == (out_dim, in_dim)


def test_reconstruct_dimension_mismatch():
    """Test error handling for dimension mismatches."""
    in_dim = 64
    out_dim = 32
    num_atoms = 16
    wrong_dim = 32
    
    dictionary = MetaDictionary(num_atoms, wrong_dim)  # Wrong dimension
    index_map = IndexMap.create_empty(out_dim, in_dim, num_atoms)
    gating_network = GatingNetwork(
        input_dim=in_dim,
        num_atoms=num_atoms,
        hidden_dim=8,
    )
    
    token_embedding = torch.randn(1, in_dim)
    
    with pytest.raises(ValueError):
        reconstruct_weights(
            dictionary, index_map, gating_network, token_embedding
        )


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

