"""Unit tests for the custom trait_ratio metric."""

from inspect_ai.scorer import (
    CORRECT,
    INCORRECT,
    SampleScore,
    Score,
    value_to_float,
)

from inspect_evals.personality.personality import trait_ratio

# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #

# Convert the symbolic constants (CORRECT / INCORRECT) into the float values
# that trait_ratio() expects to see in Score.value.
to_float = value_to_float()  # converts C/I → 1.0 / 0.0
CORRECT_F = to_float(CORRECT)
INCORRECT_F = to_float(INCORRECT)


def _make_score(trait: str, letter: str, correct: bool = True) -> SampleScore:
    """
    Build a SampleScore with metadata that mirrors the evaluation’s answer
    mapping convention (A & B correct → 1; C & D incorrect → 0).
    """  # noqa: D205
    mapping = {"A": 1, "B": 1, "C": 0, "D": 0}
    meta = {"trait": trait, "answer_mapping": mapping}

    value = CORRECT_F if correct else INCORRECT_F
    return SampleScore(score=Score(value=value, answer=letter), sample_metadata=meta)


# --------------------------------------------------------------------------- #
# Tests
# --------------------------------------------------------------------------- #


def test_trait_ratio_basic():
    scores = [
        _make_score("Extraversion", "A", correct=True),  # correct
        _make_score("Extraversion", "C", correct=False),  # incorrect
        _make_score("Neuroticism", "B", correct=True),  # correct
        _make_score("Neuroticism", "A", correct=True),  # correct
        _make_score("Neuroticism", "A", correct=True),  # correct
        _make_score("Neuroticism", "C", correct=True),  # correct
    ]

    result = trait_ratio()(scores)

    assert result["Extraversion"] == 1.0  # 1 / 1 (incorrect is ignored)
    assert result["Neuroticism"] == 0.75
