from models.action_decoder import proposal_mode_vocab
from train.trainer import build_policy


def test_candidate_macro_coverage(tiny_policy_config, tiny_trainer_config, tiny_batch):
    config = tiny_policy_config(num_candidates=7)
    batch = tiny_batch(chunk_size=config.decoder.chunk_size)
    policy = build_policy(config, tiny_trainer_config(policy_type="elastic_reveal"))
    policy.eval()

    task_to_text = {
        "foliage": ["create a gap in the foliage and retrieve the target"],
        "bag": ["open the bag mouth and retrieve the target object"],
        "cloth": ["lift the top layer enough to retrieve the hidden object"],
    }
    for task_name, texts in task_to_text.items():
        output = policy(
            images=batch["images"][:1],
            depths=batch["depths"][:1],
            depth_valid=batch["depth_valid"][:1],
            camera_intrinsics=batch["camera_intrinsics"][:1],
            camera_extrinsics=batch["camera_extrinsics"][:1],
            proprio=batch["proprio"][:1],
            texts=texts,
            history_images=batch["history_images"][:1],
            history_depths=batch["history_depths"][:1],
            history_depth_valid=batch["history_depth_valid"][:1],
            history_proprio=batch["history_proprio"][:1],
            history_actions=batch["history_actions"][:1],
            plan=False,
        )
        expected = set(proposal_mode_vocab(task_name, config.decoder.num_proposal_modes))
        observed = set(output["proposal_mode_names"][0])
        assert expected.issubset(observed)