import torch

from eval.run_reveal_benchmark import _policy_outputs
from train.trainer import build_policy


def _options(**overrides):
    base = {
        "disable_planner": False,
        "disable_memory": False,
        "disable_task_conditioning": False,
        "disable_geometry": False,
        "disable_camera_pose": False,
        "disable_shortlist": False,
        "ignore_proposal_logits": False,
        "ignore_proposal_logits_in_shortlist": False,
        "ignore_proposal_logits_in_planner": False,
        "disable_support_mode_conditioning": False,
        "disable_world_model": False,
        "disable_depth": False,
        "disable_role_tokens": False,
        "short_history": False,
        "rollout_mode_override": None,
        "use_proposal_candidates": True,
    }
    base.update(overrides)
    return base


def test_eval_toggle_paths_work(tiny_policy_config, tiny_trainer_config, tiny_batch):
    config = tiny_policy_config()
    batch = tiny_batch(chunk_size=config.decoder.chunk_size)
    model = build_policy(config, tiny_trainer_config(policy_type="elastic_reveal"))
    model.eval()

    base = _policy_outputs(model, batch, _options(), disable_planner_for_selection=False)
    no_planner = _policy_outputs(model, batch, _options(disable_planner=True), disable_planner_for_selection=False)
    no_memory = _policy_outputs(model, batch, _options(disable_memory=True), disable_planner_for_selection=True)
    no_task = _policy_outputs(model, batch, _options(disable_task_conditioning=True), disable_planner_for_selection=True)
    no_geometry = _policy_outputs(model, batch, _options(disable_geometry=True, disable_camera_pose=True), disable_planner_for_selection=True)
    no_shortlist = _policy_outputs(model, batch, _options(disable_shortlist=True), disable_planner_for_selection=False)
    ignore_shortlist_logits = _policy_outputs(
        model,
        batch,
        _options(ignore_proposal_logits_in_shortlist=True),
        disable_planner_for_selection=False,
    )
    ignore_planner_logits = _policy_outputs(
        model,
        batch,
        _options(ignore_proposal_logits_in_planner=True),
        disable_planner_for_selection=False,
    )
    ignore_proposal_logits = _policy_outputs(model, batch, _options(ignore_proposal_logits=True), disable_planner_for_selection=False)

    assert torch.count_nonzero(no_planner["planner_scores"]) == 0
    assert torch.allclose(
        no_memory["memory_output"]["memory_tokens"],
        torch.zeros_like(no_memory["memory_output"]["memory_tokens"]),
    )
    assert not torch.allclose(base["interaction_state"]["support_mode_logits"], no_task["interaction_state"]["support_mode_logits"])
    assert not torch.allclose(base["scene_tokens"], no_geometry["scene_tokens"])
    assert no_shortlist["planner_topk_indices"].shape[1] == base["candidate_chunks"].shape[1]
    assert base["proposal_logits_used_for_shortlist"]
    assert base["proposal_logits_used_for_planner"]
    assert not ignore_shortlist_logits["proposal_logits_used_for_shortlist"]
    assert ignore_shortlist_logits["proposal_logits_used_for_planner"]
    assert ignore_planner_logits["proposal_logits_used_for_shortlist"]
    assert not ignore_planner_logits["proposal_logits_used_for_planner"]
    assert torch.equal(base["planner_topk_indices"], ignore_planner_logits["planner_topk_indices"])
    assert not torch.allclose(base["planner_scores"], ignore_planner_logits["planner_scores"])
    assert not ignore_proposal_logits["proposal_logits_used_for_shortlist"]
    assert not ignore_proposal_logits["proposal_logits_used_for_planner"]
    assert not torch.allclose(base["planner_scores"], ignore_proposal_logits["planner_scores"])