import torch
from spectral_world_models.models import SelectiveSpectralTransition

def test_selective_caps_only_expansive_modes():
    m=SelectiveSpectralTransition(16,5,n_ops=2,rank=4,max_gain=1.05)
    ids=torch.arange(5)
    _, raw, clipped=m.selective_matrix(ids)
    assert float(clipped.max()) <= 1.05001
    mask=raw <= 1.05
    assert torch.allclose(clipped[mask], raw[mask], atol=1e-6)

def test_selective_forward_shape():
    m=SelectiveSpectralTransition(16,5,n_ops=2,rank=4,max_gain=1.05)
    y=m(torch.randn(3,16),torch.tensor([0,1,4]))
    assert y.shape==(3,16)