"""Utils for the ScreenQA benchmark. Referenced from https://github.com/google-research-datasets/screen_qa/blob/main/code/metrics.py"""

import collections
import re
import string

NO_ANSWER = "<no answer>"


def normalize_squad(text: str) -> str:
    """Lower text and remove punctuation, articles and extra whitespace."""

    def remove_articles(s):
        return re.sub(r"\b(a|an|the)\b", " ", s)

    def replace_punctuation(s, punc_repl):
        return "".join(punc_repl if ch in string.punctuation else ch for ch in s)

    def white_space_fix(s):
        return " ".join(s.split())

    text = text.lower()
    text = replace_punctuation(text, punc_repl="")
    text = remove_articles(text)
    text = white_space_fix(text)
    return text


def f1_score(prediction_tokens: list[str], ground_truth_tokens: list[str]) -> float:
    """Computes the F1 score between prediction and ground_truth.

    Args:
      prediction_tokens: A list of tokens in the prediction.
      ground_truth_tokens: A list of tokens in the ground truth.

    Returns:
      The F1 score between prediction and ground_truth.
    """
    common = collections.Counter(prediction_tokens) & collections.Counter(ground_truth_tokens)
    num_same = sum(common.values())
    if num_same == 0:
        return 0
    precision = 1.0 * num_same / len(prediction_tokens)
    recall = 1.0 * num_same / len(ground_truth_tokens)
    f1 = (2 * precision * recall) / (precision + recall)
    return f1


def sqa_s_metrics(prediction: str, ground_truths: list[str]) -> tuple[int, float]:
    """Computes SQA-S metrics for a single prediction.

    Args:
      prediction: The model prediction.
      ground_truths: The list of ground truth answers.

    Returns:
      A tuple of (Exact Match, F1) metrics after SQuAD preprocessing.
    """
    if prediction == NO_ANSWER:
        if any(gt == NO_ANSWER for gt in ground_truths):
            return 1, 1
        else:
            return 0, 0
    ground_truths = [gt for gt in ground_truths if gt != NO_ANSWER]
    if not ground_truths:
        return 0, 0
    prediction = normalize_squad(prediction)
    ground_truths = [normalize_squad(gt) for gt in ground_truths]
    prediction_tokens = prediction.split()
    exact_match = 1 if prediction in ground_truths else 0
    f1 = max([f1_score(prediction_tokens, gt.split()) for gt in ground_truths])
    return exact_match, f1
