from pathlib import Path
import torch
from spectral_world_models.data import DatasetConfig, generate_benchmark_npz, TransitionDataset
from spectral_world_models.models import build_model


def test_swm_forward(tmp_path: Path):
    data = tmp_path / "tiny.npz"
    generate_benchmark_npz(data, DatasetConfig(train_sequences=2, val_sequences=1, test_sequences=1, seq_len=4))
    ds = TransitionDataset(data, "train")
    batch = ds[0]
    model = build_model("swm")
    out = model(batch["image"].unsqueeze(0), batch["tokens"].unsqueeze(0), batch["action"].unsqueeze(0))
    assert out["image_pred"].shape == (1, 1, 32, 32)
    assert out["text_logits"].shape[1] == 5
    assert torch.isfinite(out["image_pred"]).all()