#!/usr/bin/env python3
"""test-set-contamination-ngram-overlap.py - is a test set in the training data? Four published overlap rules.

What it does
  Measures, for every test item, how much of it appears in a corpus of training documents, and applies the four
  overlap rules that model reports have published:
    gpt3    dirty if any N-gram of the item occurs in any training document. Words are lowercase and whitespace
            delimited with punctuation removed; N is the 5th-percentile item length in words, clipped to 8..13
            (Brown et al. 2020, arXiv 2005.14165v4, Appendix C).
    palm    contaminated if at least 70% of the item's 8-grams occur in the training data
            (Chowdhery et al. 2022, arXiv 2204.02311v5, Section 8).
    llama2  the share of the item's tokens that sit inside a run of at least 10 tokens shared with one training
            document; "dirty" at 80% or more, "clean" under 20% (Touvron et al. 2023, arXiv 2307.09288v2,
            Appendix A.6). Here the tokens are the normalised words above, not Llama 2's tokenizer, and the
            paper's four-token skip-gram budget is not applied, so matches must be exact.
    gpt4    three random 50-character substrings of the item, after removing everything but letters and digits;
            a match if any of them is a substring of one processed training example (OpenAI 2023,
            arXiv 2303.08774v6, Appendix C). Instead of drawing at random, the listing counts every 50-character
            window and reports the exact probability that three distinct random windows include a match.
  One pass gives all four: for each word position of the item, the longest run starting there that also occurs
  in some training document (found from shared 5-grams, then extended word by word).

Inputs
  --test FILE --corpus FILE  JSONL with a text field (--field, default "question") or plain text, one document
                             per line. Writes one CSV row per test item with --csv.
  --reproduce                rebuilds the article's tables: GSM8K (openai/grade-school-math at a pinned commit, MIT)
                             against its own training split and against GSM-Plus (Hugging Face qintongli/GSM-Plus at
                             a pinned revision, CC BY-SA 4.0), downloaded once into --cache and checked by SHA-256.
                             GSM-Plus is used only as text to search; nothing is trained on it.
  --profile                  with --reproduce: time and memory (tracemalloc) of indexing the GSM8K training split.

Run
  python test-set-contamination-ngram-overlap.py --reproduce --cache data --csv test-set-contamination-ngram-overlap.csv
  python test-set-contamination-ngram-overlap.py --test my_eval.jsonl --corpus my_training.jsonl --csv report.csv

Standard library only; checked on Python 3.10 and 3.13. The index is held in memory: --profile shows its cost per word.
"""
import argparse
import csv
import hashlib
import json
import math
import os
import sys
import time
import tracemalloc
import unicodedata
import urllib.request

K = 5            # candidate runs start from a shared 5-gram, so runs shorter than 5 words are reported as 0
SUB = 50         # GPT-4's substring length, in characters
GSM8K = "https://raw.githubusercontent.com/openai/grade-school-math/3101c7d5072418e28b9008a6636bde82a006892c/grade_school_math/data/"
PLUS = "https://huggingface.co/datasets/qintongli/GSM-Plus/resolve/3b708db57b96a16e8e3368ed2956990c0809440e/data/"
FILES = {"gsm8k_test.jsonl": (GSM8K + "test.jsonl", "3730d312f6e3440559ace48831e51066acaca737f6eabec99bccb9e4b3c39d14"),
         "gsm8k_train.jsonl": (GSM8K + "train.jsonl", "17f347dc51477c50d4efb83959dbb7c56297aba886e5544ee2aaed3024813465"),
         "gsmplus_test.jsonl": (PLUS + "test-00000-of-00001.jsonl", "885c279bfa9cde767ea23f79fd462d566fa0dbf6af9f2dcf30c75aa94f3d4041")}
TYPES = ["numerical substitution", "digit expansion", "integer-decimal-fraction conversion", "adding operation",
         "reversing operation", "problem understanding", "distraction insertion", "critical thinking"]
SWEEP = [5, 6, 8, 10, 13, 20, 30, 50]
COLUMNS = ["test_index", "corpus", "words", "longest_run", "ngrams8", "ngrams8_found", "tokens_covered_L10",
           "sub50_windows", "sub50_hits"]


def words(text):
    """GPT-3's gram: a lowercase, whitespace-delimited word with no punctuation (Unicode P* and S* removed)."""
    out = []
    for w in text.split():
        w = "".join(ch for ch in w if not unicodedata.category(ch).startswith(("P", "S"))).lower()
        if w:
            out.append(w)
    return out


def alnum(text):
    """GPT-4's processing: all spaces and symbols removed, letters and digits kept (case is not mentioned; kept)."""
    return "".join(ch for ch in text if ch.isalnum())


class Corpus:
    """Training documents, indexed by 5-gram position and by every 50-character window."""

    def __init__(self, docs):
        self.docs = [words(d) for d in docs]
        self.pos = {}
        for d, toks in enumerate(self.docs):
            for j in range(len(toks) - K + 1):
                self.pos.setdefault(" ".join(toks[j:j + K]), []).append((d, j))
        self.short = []                  # every processed document, for test items shorter than SUB
        self.windows = {}                # window -> the one document holding it, or -1 for several
        for d, text in enumerate(docs):
            s = alnum(text)
            self.short.append(s)
            for k in range(len(s) - SUB + 1):
                w = s[k:k + SUB]
                if self.windows.get(w, d) != d:
                    self.windows[w] = -1
                else:
                    self.windows[w] = d

    def runs(self, toks, exclude=None):
        """R[i] = the longest run starting at word i that also occurs in one document (0 if under K words)."""
        R = [0] * len(toks)
        for i in range(len(toks) - K + 1):
            best = 0
            for d, j in self.pos.get(" ".join(toks[i:i + K]), ()):
                if d == exclude:
                    continue
                doc, r = self.docs[d], K
                while i + r < len(toks) and j + r < len(doc) and toks[i + r] == doc[j + r]:
                    r += 1
                if r > best:
                    best = r
                    if i + best == len(toks):
                        break
            R[i] = best
        return R

    def substrings(self, text, exclude=None):
        """(number of 50-character windows, how many occur in some training example)."""
        s = alnum(text)
        if len(s) < SUB:                 # "or use the entire example if it's less than 50 characters"
            hit = any(s in t for d, t in enumerate(self.short) if d != exclude)
            return 1, int(hit)
        wins = [s[k:k + SUB] for k in range(len(s) - SUB + 1)]
        hits = 0
        for w in wins:
            v = self.windows.get(w)
            if v is not None and v != exclude:
                hits += 1
        return len(wins), hits


def measure(text, corpus, exclude=None):
    """One CSV row's numbers for one test item against one corpus."""
    toks = words(text)
    R = corpus.runs(toks, exclude)
    covered = set()
    for i, r in enumerate(R):
        if r >= 10:
            covered.update(range(i, i + r))
    m, h = corpus.substrings(text, exclude)
    return {"words": len(toks), "longest_run": max(R, default=0), "ngrams8": max(0, len(toks) - 7),
            "ngrams8_found": sum(1 for i in range(len(toks) - 7) if R[i] >= 8),
            "tokens_covered_L10": len(covered), "sub50_windows": m, "sub50_hits": h}


def gpt3_n(word_counts):
    """N = the 5th-percentile item length in words (nearest rank), clipped to 8..13."""
    s = sorted(word_counts)
    p5 = s[math.ceil(0.05 * len(s)) - 1]
    return p5, min(13, max(8, p5))


def digit_windows(texts, n):
    """(word windows of length n that contain a digit, all word windows of length n), pooled over texts."""
    hit = total = 0
    for t in texts:
        w = words(t)
        has = [any(ch.isdigit() for ch in x) for x in w]
        for i in range(len(w) - n + 1):
            total += 1
            hit += any(has[i:i + n])
    return hit, total


def p_three(m, h):
    """Probability that three distinct random windows out of m include at least one of the h matching ones."""
    if m < 3:
        return 1.0 if h else 0.0
    return 1 - math.comb(m - h, 3) / math.comb(m, 3)


def flags(row, n):
    """The four rules' verdicts for one CSV row (numbers as written to the CSV)."""
    w, g, cov = int(row["words"]), int(row["ngrams8"]), int(row["tokens_covered_L10"])
    return {"gpt3": int(row["longest_run"]) >= n,
            "palm": g > 0 and 10 * int(row["ngrams8_found"]) >= 7 * g,       # at least 70%, in integers
            "llama2_dirty": w > 0 and 10 * cov >= 8 * w,                     # 80% or more
            "llama2_not_clean": w > 0 and 10 * cov >= 2 * w,                 # 20% or more ("clean" is under 20%)
            "gpt4_p3": p_three(int(row["sub50_windows"]), int(row["sub50_hits"])),
            "gpt4_any": int(row["sub50_hits"]) > 0}


def read_texts(path, field):
    out = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.rstrip("\n")
            if not line.strip():
                continue
            if line.lstrip().startswith("{"):
                out.append(json.loads(line)[field])
            else:
                out.append(line)
    return out


def fetch(cache):
    os.makedirs(cache, exist_ok=True)
    for name, (url, sha) in FILES.items():
        path = os.path.join(cache, name)
        if not os.path.exists(path):
            req = urllib.request.Request(url, headers={"User-Agent": "test-set-contamination-ngram-overlap.py"})
            with urllib.request.urlopen(req, timeout=120) as r:
                data = r.read()
            with open(path, "wb") as f:
                f.write(data)
        with open(path, "rb") as f:
            got = hashlib.sha256(f.read()).hexdigest()
        if got != sha:
            sys.exit(f"{name}: SHA-256 {got} is not the pinned {sha}")


def pct(k, n):
    return f"{100 * k / n:.1f}%"


def reproduce(cache, csv_path, profile):
    fetch(cache)
    test = read_texts(os.path.join(cache, "gsm8k_test.jsonl"), "question")
    train = read_texts(os.path.join(cache, "gsm8k_train.jsonl"), "question")
    plus = [json.loads(line) for line in open(os.path.join(cache, "gsmplus_test.jsonl"), encoding="utf-8")]
    index = {q: i for i, q in enumerate(test)}
    assert len(index) == len(test) == 1319 and len(train) == 7473 and len(plus) == 10552
    variant = {}
    for p in plus:
        variant[(index[p["seed_question"]], p["perturbation_type"])] = p["question"]
    assert len(variant) == 1319 * 8 and {t for _, t in variant} == set(TYPES)
    p5, n = gpt3_n([len(words(q)) for q in test])
    print(f"GSM8K test: {len(test)} questions; GSM8K train: {len(train)}; GSM-Plus: {len(plus)} variations")
    print(f"5th-percentile test question: {p5} words, so GPT-3's N = {n}")
    for w in (8, n):
        hit, total = digit_windows(test, w)
        print(f"{w}-word windows of the test questions that contain a digit: {hit:,} of {total:,} ({pct(hit, total)})")

    if profile:
        tracemalloc.start()
        t0 = time.perf_counter()
        c = Corpus(train)
        dt = time.perf_counter() - t0
        size, _ = tracemalloc.get_traced_memory()
        tracemalloc.stop()
        positions = sum(len(d) for d in c.docs)
        print(f"profile: {positions} word positions, {len(c.pos)} distinct 5-grams, {len(c.windows)} distinct "
              f"50-char windows; index {size / 2**20:.1f} MiB = {size / positions:.0f} bytes per word position; "
              f"built in {dt:.2f} s")
        return

    rows = []
    corpora = [("gsm8k_train", Corpus(train), False), ("gsm8k_test_others", Corpus(test), True),
               ("gsmplus_all", Corpus([p["question"] for p in plus]), False)]
    for name, corpus, self_excl in corpora:
        for i, q in enumerate(test):
            rows.append({"test_index": i, "corpus": name, **measure(q, corpus, i if self_excl else None)})
    for t in TYPES:
        for i, q in enumerate(test):
            rows.append({"test_index": i, "corpus": "own_variant:" + t, **measure(q, Corpus([variant[(i, t)]]))})
    order = ["gsm8k_train", "gsm8k_test_others"] + ["own_variant:" + t for t in TYPES] + ["gsmplus_all"]
    rows.sort(key=lambda r: (order.index(r["corpus"]), r["test_index"]))
    if csv_path:
        with open(csv_path, "w", newline="", encoding="utf-8") as f:
            w = csv.DictWriter(f, fieldnames=COLUMNS, lineterminator="\n")
            w.writeheader()
            w.writerows(rows)

    by = {}
    for r in rows:
        by.setdefault(r["corpus"], []).append({k: str(v) for k, v in r.items()})
    names = {"gsm8k_train": "GSM8K train (7,473 other problems)", "gsm8k_test_others": "GSM8K test, the other questions",
             "gsmplus_all": "all 10,552 GSM-Plus variations"}

    def label(name):
        return names[name] if name in names else "own variant: " + name.split(":", 1)[1]
    print()
    print(f"Share of the 1,319 test questions each rule flags (gpt3 at N = {n})")
    print(f"{'training corpus':48s} {'gpt3':>6s} {'palm':>6s} {'l2 dirty':>9s} {'l2 >=20%':>9s} {'gpt4 x3':>8s} {'gpt4 all':>9s}")
    for name in order:
        fl = [flags(r, n) for r in by[name]]
        cnt = len(fl)
        row = [pct(sum(f["gpt3"] for f in fl), cnt), pct(sum(f["palm"] for f in fl), cnt),
               pct(sum(f["llama2_dirty"] for f in fl), cnt), pct(sum(f["llama2_not_clean"] for f in fl), cnt),
               f"{100 * sum(f['gpt4_p3'] for f in fl) / cnt:.1f}%", pct(sum(f["gpt4_any"] for f in fl), cnt)]
        nm = label(name)
        print(f"{nm:48s} {row[0]:>6s} {row[1]:>6s} {row[2]:>9s} {row[3]:>9s} {row[4]:>8s} {row[5]:>9s}")
    print()
    print("Share flagged by a single shared run of at least n words (any-hit at n)")
    print(f"{'training corpus':48s} " + " ".join(f"{'n=' + str(s):>6s}" for s in SWEEP))
    for name in order:
        runs = [int(r["longest_run"]) for r in by[name]]
        nm = label(name)
        print(f"{nm:48s} " + " ".join(f"{pct(sum(x >= s for x in runs), len(runs)):>6s}" for s in SWEEP))
    print()
    for name, what in (("gsm8k_train", "a training question"), ("gsm8k_test_others", "another test question")):
        hits = sorted(((int(r["longest_run"]), int(r["test_index"])) for r in by[name] if int(r["longest_run"]) >= n),
                      key=lambda x: (-x[0], x[1]))
        print(f"Test items sharing a run of {n}+ words with {what}: {len(hits)} "
              f"({'; '.join(f'item {i}: {k} words' for k, i in hits)})")


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--test")
    ap.add_argument("--corpus")
    ap.add_argument("--field", default="question")
    ap.add_argument("--csv")
    ap.add_argument("--reproduce", action="store_true")
    ap.add_argument("--cache", default="data")
    ap.add_argument("--profile", action="store_true")
    a = ap.parse_args()
    sys.stdout.reconfigure(encoding="utf-8")
    if a.reproduce:
        return reproduce(a.cache, a.csv, a.profile)
    if not (a.test and a.corpus):
        ap.error("give --test and --corpus, or --reproduce")
    test, docs = read_texts(a.test, a.field), read_texts(a.corpus, a.field)
    if not test or not docs:
        sys.exit("no test items or no training documents read")
    p5, n = gpt3_n([len(words(q)) for q in test])
    corpus = Corpus(docs)
    rows = [{"test_index": i, "corpus": os.path.basename(a.corpus), **measure(q, corpus)} for i, q in enumerate(test)]
    fl = [flags({k: str(v) for k, v in r.items()}, n) for r in rows]
    print(f"{len(test)} test items, {len(docs)} training documents; GPT-3 N = {n} (5th percentile {p5} words)")
    for key in ("gpt3", "palm", "llama2_dirty", "llama2_not_clean", "gpt4_any"):
        print(f"  {key:17s} {sum(f[key] for f in fl):6d}  {pct(sum(f[key] for f in fl), len(fl))}")
    expected = sum(f["gpt4_p3"] for f in fl)
    print(f"  {'gpt4_p3':17s} {expected:6.1f}  {100 * expected / len(fl):.1f}%  (expected, three random substrings)")
    if a.csv:
        with open(a.csv, "w", newline="", encoding="utf-8") as f:
            w = csv.DictWriter(f, fieldnames=COLUMNS, lineterminator="\n")
            w.writeheader()
            w.writerows(rows)


if __name__ == "__main__":
    main()