from inspect_ai.model import ChatCompletionChoice, ChatMessageAssistant, ModelOutput
from inspect_ai.scorer import CORRECT, INCORRECT, NOANSWER, Score, Scorer, Target
from inspect_ai.solver import TaskState

from inspect_evals.lab_bench.record_to_sample_helpers import UNCERTAIN_ANSWER_CHOICE
from inspect_evals.lab_bench.scorer import precision_choice

scorer: Scorer = precision_choice(no_answer=UNCERTAIN_ANSWER_CHOICE)


model = "gpt-4o-mini-2024-07-18"
sample_id = 1
epoch = 1
input = "What effect does expression of the ATPase-deficient E251Q mutant of the Spindle E protein in silkworm cells have on the levels of the mature piwiRNAs piR1712 and piR2986?"
messages: list = []


async def test_precision_choice_correct():
    task_state_instance = TaskState(
        model,
        sample_id,
        epoch,
        input,
        messages,
        choices=[
            "Decrease",
            "Increase",
            "No effect",
            UNCERTAIN_ANSWER_CHOICE,
        ],
        output=ModelOutput(
            choices=[
                ChatCompletionChoice(message=ChatMessageAssistant(content="ANSWER: A"))
            ]
        ),
    )
    task_state_instance.choices.mark_choice(0, True)
    target_instance = Target("A")
    result = await scorer(task_state_instance, target_instance)
    assert result == Score(value=CORRECT, answer="A", explanation="ANSWER: A")


async def test_precision_choice_incorrect():
    task_state_instance = TaskState(
        model,
        sample_id,
        epoch,
        input,
        messages,
        choices=[
            "Decrease",
            "Increase",
            "No effect",
            UNCERTAIN_ANSWER_CHOICE,
        ],
        output=ModelOutput(
            choices=[
                ChatCompletionChoice(message=ChatMessageAssistant(content="ANSWER: B"))
            ]
        ),
    )
    task_state_instance.choices.mark_choice(1, True)
    target_instance = Target("A")
    result = await scorer(task_state_instance, target_instance)
    assert result == Score(value=INCORRECT, answer="B", explanation="ANSWER: B")


async def test_precision_choice_no_answer():
    task_state_instance = TaskState(
        model,
        sample_id,
        epoch,
        input,
        messages,
        choices=[
            "Decrease",
            "Increase",
            "No effect",
            UNCERTAIN_ANSWER_CHOICE,
        ],
        output=ModelOutput(
            choices=[
                ChatCompletionChoice(message=ChatMessageAssistant(content="ANSWER: D"))
            ]
        ),
    )
    task_state_instance.choices.mark_choice(3, True)
    target_instance = Target("A")
    result = await scorer(task_state_instance, target_instance)
    assert result == Score(value=NOANSWER, answer="D", explanation="ANSWER: D")
