"""Evaluation-set decontamination by shared token n-grams."""

from __future__ import annotations

from collections.abc import Iterable


def build_ngram_index(token_lists: Iterable[list[int]], n: int = 13, letter_ids: set[int] | None = None) -> set[int]:
    """Hashes of every n-gram in the eval texts.

    With `letter_ids`, n-grams containing no token with a letter are skipped: runs of digits and
    punctuation (e.g. 13 zeros of a memory address) are shared by unrelated documents and do not
    indicate leaked evaluation text.
    """
    index: set[int] = set()
    for ids in token_lists:
        for i in range(len(ids) - n + 1):
            gram = tuple(ids[i : i + n])
            if letter_ids is None or not letter_ids.isdisjoint(gram):
                index.add(hash(gram))
    return index


def is_contaminated(ids: list[int], index: set[int], n: int = 13) -> bool:
    if not index:
        return False
    return any(hash(tuple(ids[i : i + n])) in index for i in range(len(ids) - n + 1))