from __future__ import annotations

import pytest

from dovla_cil.transfercritic.eval import compare_selection_strategies
from dovla_cil.transfercritic.labeling import make_utility_labels, toy_utility_value
from dovla_cil.transfercritic.schema import DataAtom, TransferContext
from dovla_cil.transfercritic.selection import greedy_marginal_selection


def _atom(
    atom_id: str,
    *,
    score: float,
    success: bool = False,
    candidate_type: str = "near_miss",
    task_id: str = "task_a",
    cost: float = 1.0,
) -> DataAtom:
    return DataAtom(
        record_id=atom_id,
        embedding=[score, 1.0 if success else 0.0, cost, 0.5],
        candidate_type=candidate_type,
        task_metadata={"task_id": task_id, "family": "pick"},
        reward_summary={
            "progress": score,
            "score": score + (1.0 if success else 0.0),
            "success": 1.0 if success else 0.0,
            "regret": max(0.0, 1.0 - score),
        },
        effect_summary={"moved_object_count": 1.0, "true_relation_count": score},
        cost=cost,
    )


def test_data_atom_and_context_schema_roundtrip() -> None:
    atom = _atom("r1", score=0.5, success=True)
    context = TransferContext(
        benchmark_name="CausalStress",
        task_family="pick",
        target_objects=["red_mug"],
        ood_factor="wrong_target",
        validation_ref="val://small",
    )

    restored = DataAtom.from_dict(atom.to_dict())
    context_restored = TransferContext.from_dict(context.to_dict())

    assert restored == atom
    assert context_restored == context
    assert restored.atom_id == "r1"


def test_greedy_selection_uses_score_over_cost() -> None:
    atoms = [
        _atom("cheap_good", score=0.8, success=True, cost=1.0),
        _atom("expensive_best", score=1.0, success=True, cost=3.0),
        _atom("bad", score=0.1, success=False, cost=1.0),
    ]
    context = TransferContext(benchmark_name="CausalStress", task_family="pick")

    result = greedy_marginal_selection(
        atoms,
        context,
        budget=2.0,
        score_fn=lambda atom, _selected, _context: float(atom.reward_summary["score"]),
    )

    assert result.name == "transfercritic"
    assert result.total_cost <= 2.0
    assert result.atom_ids[0] == "cheap_good"
    assert "expensive_best" not in result.atom_ids


def test_toy_utility_labels_prefer_successful_useful_atoms() -> None:
    context = TransferContext(benchmark_name="CausalStress", task_family="pick")
    good = _atom("good", score=0.9, success=True, candidate_type="near_miss")
    poor = _atom("poor", score=0.1, success=False, candidate_type="noop")

    labels = make_utility_labels([good, poor], context, method="toy_retraining_delta")

    assert labels[0].utility == toy_utility_value(good, context)
    assert labels[0].utility > labels[1].utility
    assert labels[0].metadata["approximate"] is True


def test_selection_experiments_return_all_baselines() -> None:
    atoms = [
        _atom("a", score=0.6, task_id="task_a"),
        _atom("b", score=0.9, success=True, task_id="task_b"),
        _atom("c", score=0.2, task_id="task_a"),
    ]
    context = TransferContext(benchmark_name="CausalStress", task_family="pick")

    rows = compare_selection_strategies(atoms, context, budget=2.0, seed=0)

    assert {row["name"] for row in rows} == {
        "random_subset",
        "top_reward_subset",
        "task_balanced_subset",
        "transfercritic",
    }
    assert all(float(row["total_cost"]) <= 2.0 for row in rows)


def test_transfercritic_model_shapes() -> None:
    torch = pytest.importorskip("torch")
    from dovla_cil.transfercritic.model import SetConditionedTransferCritic, TransferCriticConfig

    config = TransferCriticConfig(atom_dim=4, set_dim=4, context_dim=4, hidden_dim=8)
    model = SetConditionedTransferCritic(config)
    atom = torch.randn(3, 4)
    current_set = torch.zeros(3, 4)
    context = torch.randn(3, 4)

    score = model(atom, current_set, context)

    assert score.shape == (3,)
    score.mean().backward()
    assert any(parameter.grad is not None for parameter in model.parameters())