tests/test_data_contract.py

from __future__ import annotations

import numpy as np
import pytest
import torch
from types import SimpleNamespace

from src.data_contract import (
    CosmosRBMCollator,
    adapt_official_trajectory,
    audit_group_overlap,
    deterministic_frame_indices,
    make_preference_contract,
    prepare_trajectory,
    validate_span_bounds,
    worker_invariant_sample_ids,
)
from src.reward_heads import SharedRewardHeads


def synthetic_frames(count: int, height: int = 12, width: int = 20) -> np.ndarray:
    frames = np.empty((count, height, width, 3), dtype=np.uint8)
    for index in range(count):
        frames[index].fill(index)
    return frames


def prepared(
    *,
    strategy: str = "forward",
    quality: str = "successful",
    frame_count: int = 13,
    trajectory_id: str = "trajectory-a",
):
    return prepare_trajectory(
        trajectory_id=trajectory_id,
        frames=synthetic_frames(frame_count),
        task="insert cable",
        data_source="unit_test",
        quality_label=quality,
        strategy=strategy,
        canvas_size=16,
    )


def test_forward_progress_is_nondecreasing_and_sampling_is_exact() -> None:
    item = prepared()
    expected, _ = deterministic_frame_indices(13, 8)
    assert np.array_equal(item.source_frame_indices, expected)
    assert np.all(np.diff(item.progress) >= 0)
    assert item.frames_uint8.shape == (8, 16, 16, 3)


def test_reverse_keeps_all_frame_fields_synchronized() -> None:
    forward = prepared()
    reverse = prepared(strategy="reverse")
    assert np.array_equal(reverse.source_frame_indices, forward.source_frame_indices[::-1])
    assert np.array_equal(reverse.progress, forward.progress[::-1])
    assert np.array_equal(reverse.frame_valid_mask, forward.frame_valid_mask[::-1])
    assert np.array_equal(reverse.success_labels, forward.success_labels[::-1])


def test_rewind_has_a_real_progress_drop() -> None:
    rewind = prepared(strategy="rewind")
    assert np.any(np.diff(rewind.progress) < 0)


def test_different_task_is_zero_with_supervision() -> None:
    item = prepared(strategy="different_task", quality="failure")
    assert np.all(item.progress == 0)
    assert np.array_equal(item.progress_loss_mask, item.frame_valid_mask)


def test_preference_swap_flips_label_and_bradley_terry_logit() -> None:
    chosen = prepared(trajectory_id="chosen")
    rejected = prepared(strategy="rewind", trajectory_id="rejected")
    pair = make_preference_contract(chosen, rejected, strategy="rewind")
    swapped = pair.swapped()
    assert swapped.preference_label == 0.0
    assert swapped.trajectory_a.trajectory_id == "rejected"

    torch.manual_seed(7)
    heads = SharedRewardHeads(hidden_size=8, dropout=0.0)
    features_a = torch.randn(1, 8, 8)
    features_b = torch.randn(1, 8, 8)
    mask = torch.ones(1, 8, dtype=torch.bool)
    ab = heads.preference_logit(features_a, mask, features_b, mask)
    ba = heads.preference_logit(features_b, mask, features_a, mask)
    assert torch.allclose(ab, -ba)


def test_padding_never_enters_loss_or_pooling() -> None:
    item = prepared(frame_count=3)
    batch = CosmosRBMCollator()([item])["trajectory_inputs"]
    assert batch["frame_valid_mask"].tolist() == [
        [True, True, True, False, False, False, False, False]
    ]
    assert not batch["progress_loss_mask"][0, 3:].any()
    assert not batch["success_loss_mask"][0, 3:].any()

    pool_features = torch.randn(1, 8, 4)
    heads = SharedRewardHeads(hidden_size=4, dropout=0.0)
    score_a = heads(pool_features, batch["frame_valid_mask"])["trajectory_scores"]
    pool_features[:, 3:] = 1000
    score_b = heads(pool_features, batch["frame_valid_mask"])["trajectory_scores"]
    assert torch.allclose(score_a, score_b)


def test_failure_does_not_receive_dense_progress() -> None:
    failure = prepared(quality="failure")
    assert not failure.progress_loss_mask.any()
    assert not failure.success_labels.any()


def test_official_sampler_adapter_preserves_labels_and_indices() -> None:
    official_progress = np.linspace(0.0, 1.0, 8, dtype=np.float32)
    official_success = [0.0] * 7 + [1.0]
    trajectory = SimpleNamespace(
        id="official-a",
        task="insert cable",
        data_source="official_unit_test",
        quality_label="successful",
        frames=synthetic_frames(8),
        target_progress=official_progress.tolist(),
        success_label=official_success,
        predict_last_frame_mask=[1.0] * 8,
        metadata={
            "selected_frame_indices": [0, 4, 9, 13, 18, 22, 27, 31],
            "frame_valid_mask": [True] * 8,
        },
    )
    adapted = adapt_official_trajectory(
        trajectory,
        strategy="forward",
        progress_supervised=True,
        canvas_size=16,
    )
    assert np.allclose(adapted.progress, official_progress)
    assert adapted.success_labels.tolist() == official_success
    assert adapted.source_frame_indices.tolist() == trajectory.metadata[
        "selected_frame_indices"
    ]
    assert adapted.metadata["adapter"] == "official_robometer_sampler_v1"


def test_all_backbones_consume_the_same_frozen_indices() -> None:
    frozen, _ = deterministic_frame_indices(32, 8)
    for _backbone in ("qwen3-vl", "cosmos3-edge", "frozen-cosmos"):
        item = prepared(frame_count=32)
        assert np.array_equal(item.source_frame_indices, frozen)


def test_visual_span_contract() -> None:
    validate_span_bounds([(1, 4), (8, 11), (12, 15)], 3, 16)
    with pytest.raises(ValueError, match="overlap"):
        validate_span_bounds([(1, 5), (4, 8)], 2, 10)
    with pytest.raises(ValueError, match="count"):
        validate_span_bounds([(1, 5)], 2, 10)


def test_worker_count_does_not_change_seeded_sampling_list() -> None:
    ids = [f"trajectory-{index}" for index in range(100)]
    expected = worker_invariant_sample_ids(ids, seed=17, worker_count=1)
    assert worker_invariant_sample_ids(ids, seed=17, worker_count=2) == expected
    assert worker_invariant_sample_ids(ids, seed=17, worker_count=8) == expected


def test_group_overlap_gate() -> None:
    assert audit_group_overlap(["a", "b", "c"], ["d", "e"]) == set()
    assert audit_group_overlap(["a", "b"], ["b", "c"]) == {"b"}


def test_one_hundred_trajectories_cpu_collate_and_backward() -> None:
    item = prepared()
    batch = CosmosRBMCollator()([item] * 100)["trajectory_inputs"]
    assert batch["frames_uint8"].shape == (100, 8, 16, 16, 3)
    features = torch.randn(8, 8, 16, requires_grad=True)
    heads = SharedRewardHeads(hidden_size=16, dropout=0.0)
    output = heads(features, batch["frame_valid_mask"][:8])
    loss = output["progress_logits"].float().mean() + output[
        "success_logits"
    ].float().mean()
    loss.backward()
    assert features.grad is not None and torch.isfinite(features.grad).all()