import torch

from stimulus_synthesis.outputs import StimulusCandidate, StimulusSynthesisOutput
from stimulus_synthesis.scoring.objectives import indices_mean, vector_cosine, vector_dot


def test_objectives():
    predictions = torch.tensor([[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]])

    assert torch.allclose(indices_mean(predictions, {"type": "indices", "indices": [0, 2]}), torch.tensor([2.0, 2.0]))
    assert torch.allclose(vector_dot(predictions, {"type": "vector", "vector": [1.0, 0.0, 0.0]}), torch.tensor([1.0, 3.0]))
    assert vector_cosine(predictions, {"type": "vector", "vector": [1.0, 0.0, 0.0]}).shape == (2,)


def test_output_schema():
    candidate = StimulusCandidate(prompt="person running", score=1.5, image="image", video="video")
    output = StimulusSynthesisOutput(
        candidates=[candidate],
        best_prompt=candidate.prompt,
        best_score=candidate.score,
        history_best=[1.5],
    )

    assert output.best.prompt == "person running"
    assert output.best_score == 1.5