from inspect_ai.scorer import (
    CORRECT,
    INCORRECT,
    NOANSWER,
    Score,
    Scorer,
    Target,
    accuracy,
    scorer,
    stderr,
)
from inspect_ai.solver import TaskState

from inspect_evals.lab_bench.metrics import coverage, precision


@scorer(metrics=[accuracy(), precision(), coverage(), stderr()])
def precision_choice(no_answer: str | None = None) -> Scorer:
    async def score(state: TaskState, target: Target) -> Score:
        choices = state.choices
        explanation = state.output.completion

        # compute the target positions
        target_positions = [
            ord(target_character.upper()) - ord("A") for target_character in target.text
        ]

        # compute the answers
        choice_positions = [
            i for i, choice in enumerate(choices) if choice.correct is True
        ]
        answers = [chr(ord("A") + choice) for choice in choice_positions]

        # the choices the model selected
        generated_selected_choices = [
            i for i, choice in enumerate(choices) if choice.correct is True
        ]

        if no_answer is not None:
            # Look for a no-answer choice and use that if present
            no_answer_choices = [
                i
                for i, choice in enumerate(choices)
                if choice.value == no_answer and choice.correct is True
            ]

            if len(no_answer_choices) == 1:
                return Score(
                    value=NOANSWER,
                    answer=", ".join(answers),
                    explanation=explanation,
                )

        target_matches_choices = generated_selected_choices == sorted(target_positions)
        return Score(
            value=CORRECT if target_matches_choices else INCORRECT,
            answer=", ".join(answers),
            explanation=explanation,
        )

    return score
