"""Cross-modal integration tests for Phase 20."""
import os
import torch
import pytest
from arbitor.main import ARBModel

pytestmark = pytest.mark.skipif(
    os.environ.get("ARB_RUN_SLOW_TESTS") != "1",
    reason="full cross-modal ARBModel tests require the 3B target model and sidecar encoders",
)


def test_cross_modality_unified_latent():
    model = ARBModel(enable_image=True, enable_audio=True)
    model.eval()
    text = torch.randint(0, 288, (1, 50))
    img = torch.randn(1, 3, 256, 256)
    audio = torch.randn(1, 16000 * 3)
    logits, losses, indices, _ = model(text, images=img, audio=audio)
    assert logits is not None
    assert indices is not None
    assert indices.shape[1] > 50


def test_text_only_still_works():
    model = ARBModel(enable_image=False, enable_audio=False)
    model.eval()
    text = torch.randint(0, 288, (1, 50))
    logits, losses, indices, _ = model(text)
    assert logits is not None
    assert logits.shape[-1] == 288


def test_image_only():
    model = ARBModel(enable_image=True, enable_audio=False)
    model.eval()
    text = torch.randint(0, 288, (1, 10))
    img = torch.randn(1, 3, 224, 224)
    logits, losses, indices, _ = model(text, images=img)
    assert logits is not None
    assert indices.shape[1] > 10