from typing import Any

import nltk
from nltk.tokenize import word_tokenize

from genderbench.probing.attempt import Attempt
from genderbench.probing.evaluator import ClosedSetEvaluator


class YesNoEvaluator(ClosedSetEvaluator):
    """
    Detect "yes" or "no" answers. If both or neither are present, the evaluation
    remains `undefined`.
    """

    def __init__(self, probe):
        super().__init__(probe=probe, options=["yes", "no"])
        nltk.download("punkt", quiet=True)

    def calculate_evaluation(self, attempt: Attempt) -> Any:
        tokens = word_tokenize(attempt.answer.lower())
        contains_yes, contains_no = "yes" in tokens, "no" in tokens
        if contains_yes == contains_no:
            return self.undetected
        if contains_yes:
            return "yes"
        return "no"
