from __future__ import annotations

from dovla_cil.interventions.samplers import InterventionSampler, RandomInterventionSampler
from dovla_cil.sim.registry import get_simulator_backend
from dovla_cil.tasks.library import ToyTaskLibrary


def test_sampler_returns_exactly_k_or_less_when_impossible() -> None:
    task = ToyTaskLibrary().get_by_id("toy_pick_object_among_distractors")
    sim = get_simulator_backend("toy")
    sim.seed(3)
    sim.reset_task(task)
    expert = RandomInterventionSampler(k=1, seed=3).sample(sim.render_observation(), task)[0].action

    actions = InterventionSampler(seed=9).sample(
        task=task,
        observation=sim.render_observation(),
        symbolic_state=sim.get_symbolic_state(),
        expert_actions=[expert],
        k=6,
    )

    assert len(actions) == 6
    assert len({action.action_id for action in actions}) == len(actions)


def test_sampler_candidate_metadata_and_determinism() -> None:
    task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl")
    sim = get_simulator_backend("toy")
    sim.seed(5)
    sim.reset_task(task)
    expert = RandomInterventionSampler(k=1, seed=5).sample(sim.render_observation(), task)[0].action

    first = InterventionSampler(seed=11).sample(
        task, sim.render_observation(), sim.get_symbolic_state(), [expert], 8
    )
    second = InterventionSampler(seed=11).sample(
        task, sim.render_observation(), sim.get_symbolic_state(), [expert], 8
    )

    assert [action.to_dict() for action in first] == [action.to_dict() for action in second]
    for action in first:
        assert "candidate_type" in action.metadata
        assert "intended_target" in action.metadata
        assert "intended_relation" in action.metadata
        assert "difficulty" in action.metadata


def test_wrong_target_actions_use_distractors_if_present() -> None:
    task = ToyTaskLibrary().get_by_id("toy_pick_object_among_distractors")
    sim = get_simulator_backend("toy")
    sim.seed(7)
    sim.reset_task(task)

    actions = InterventionSampler(seed=7).sample(
        task,
        sim.render_observation(),
        sim.get_symbolic_state(),
        [],
        8,
    )
    wrong_targets = [
        action.metadata["intended_target"]
        for action in actions
        if action.metadata.get("candidate_type") == "wrong_target"
    ]

    assert wrong_targets
    assert set(wrong_targets).issubset(set(task.distractor_object_ids))


def test_near_miss_actions_differ_from_expert() -> None:
    task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl")
    sim = get_simulator_backend("toy")
    sim.seed(13)
    sim.reset_task(task)
    expert = RandomInterventionSampler(k=1, seed=13).sample(
        sim.render_observation(), task
    )[0].action

    actions = InterventionSampler(seed=13).sample(
        task,
        sim.render_observation(),
        sim.get_symbolic_state(),
        [expert],
        8,
    )
    near_misses = [
        action for action in actions if action.metadata.get("candidate_type") == "near_miss"
    ]

    assert near_misses
    assert all(action.values != expert.values for action in near_misses)