tests/test_system1.py

"""System-1 multimodal decision model tests (tiny backend, CPU-only)."""

import sys
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

torch = pytest.importorskip("torch")

from qwenjev.candidates import CandidateAction  # noqa: E402
from qwenjev.system1 import QwenJevModel, System1Config  # noqa: E402
from qwenjev.vision import TinyVLMConfig  # noqa: E402

GRID = [
    [1, 1, 1, 1, 1],
    [1, 2, 0, 3, 1],
    [1, 0, 4, 0, 1],
    [1, 1, 1, 1, 1],
]

CANDS = [
    CandidateAction("ACTION1", text="move up"),
    CandidateAction("ACTION3", text="move left"),
    CandidateAction("ACTION6(2,2)", x=2, y=2, text="click (2,2)"),
]


@pytest.fixture(scope="module")
def model():
    cfg = System1Config(
        backend="tiny",
        tiny=TinyVLMConfig(image_size=64, patch_size=8, lm_dim=64,
                           vision_dim=48, max_text_len=256),
    )
    return QwenJevModel(cfg).build()


def test_choice_probabilities_normalized(model):
    d = model.choose("qwenjev-state-v1|H=4 W=5 L=0 S=1", CANDS, grid=GRID)
    assert len(d.probs) == len(CANDS)
    assert abs(sum(d.probs) - 1.0) < 1e-5
    assert all(p >= 0 for p in d.probs)
    assert d.chosen in CANDS


def test_choice_requires_two_candidates(model):
    with pytest.raises(ValueError):
        model.choose("state", CANDS[:1], grid=GRID)


def test_vision_changes_choice_scores(model):
    """Regression: the frame image MUST influence the decision scores."""
    torch.manual_seed(0)
    state = "qwenjev-state-v1|H=4 W=5 L=0 S=1"
    g1 = [[v for v in row] for row in GRID]
    g2 = [[v for v in row] for row in GRID]
    g2[1][2] = 9                      # perturb one cell
    leaves1 = model.encode_question(state, CANDS, grid=g1)
    leaves2 = model.encode_question(state, CANDS, grid=g2)
    p1 = model.forward_choice(leaves1)
    p2 = model.forward_choice(leaves2)
    assert max(abs(a - b) for a, b in zip(p1, p2)) > 1e-6, (
        "vision tower has no effect on Choice scores"
    )


def test_vision_off_ablation(model):
    """use_vision=False still works (ablation mode) and stays finite."""
    old = model.config.use_vision
    try:
        model.config.use_vision = False
        leaves = model.encode_question("state", CANDS, grid=None)
        for leaf in leaves:
            assert leaf.pixel_values is None
        probs = model.forward_choice(leaves)
        assert abs(sum(probs) - 1.0) < 1e-5
    finally:
        model.config.use_vision = old


def test_boolean_and_score_heads(model):
    state = "qwenjev-state-v1|H=4 W=5 L=0 S=1"
    p = model.forward_boolean(state, "inst", "ACTION3 changes the frame", GRID)
    assert 0.0 <= p <= 1.0
    s = model.forward_score(state, "inst", "promise of state", GRID)
    assert 2.0 <= s <= 10.0


def test_boolean_logits_are_trainable(model):
    """Trainable twin must keep the autograd graph (old bug: detached)."""
    logits = model.forward_boolean_logits("state", "inst", "prop", GRID)
    assert logits.requires_grad
    logits.sum().backward()           # no "does not require grad" error


def test_score_logits_are_trainable(model):
    logits = model.forward_score_logits("state", "inst", "desc", GRID)
    assert logits.requires_grad


def test_save_and_load_heads(model, tmp_path):
    out = tmp_path / "ckpt"
    model.save(str(out))
    assert (out / "qwenjev_s1.pt").exists()

    cfg = System1Config(
        backend="tiny",
        tiny=TinyVLMConfig(image_size=64, patch_size=8, lm_dim=64,
                           vision_dim=48, max_text_len=256),
    )
    m2 = QwenJevModel(cfg).build()
    m2.load_heads(str(out))
    model.eval(), m2.eval()               # deterministic comparison
    state = "qwenjev-state-v1|H=4 W=5 L=0 S=1"
    p1 = model.forward_choice(model.encode_question(state, CANDS, grid=GRID))
    p2 = m2.forward_choice(m2.encode_question(state, CANDS, grid=GRID))
    assert max(abs(a - b) for a, b in zip(p1, p2)) < 1e-5


def test_hf_backend_config_smoke():
    """hf backend is the real-model path; verify config plumbing only
    (no weights download in CI)."""
    cfg = System1Config(backend="hf", hf_model="Qwen/Qwen3-VL-2B-Instruct")
    assert cfg.hf_model == "Qwen/Qwen3-VL-2B-Instruct"
    assert cfg.freeze_vision_tower is False