"""Count tune items whose text the tuner copied into the final ungated skill (3.9 "it is memorising the questions").

The names and terms in VOCAB are first removed from both the item and the skill, so a rule that merely spells a term
right ("test a LoRA with MLX") does not count. An item then counts as copied if any N consecutive tokens of its raw
transcript or its ground truth appear in the skill. Tokens are lowercase words for en (N = 8) and characters for zh
(N = 10), punctuation and whitespace removed, so a copied sentence still matches after the tuner changes quotes or
line breaks. Unseen items (holdout + test) are checked the same way as a control: the tuner never saw them, so hits
there would mean the window is too short.
Usage: copy_check.py <lang> [--round 8] [--n N]"""
import json, re, sys
from pathlib import Path

from lib import VOCAB

HERE = Path(__file__).resolve().parent
lang = sys.argv[1]
rnd = int(sys.argv[sys.argv.index("--round") + 1]) if "--round" in sys.argv else 8
N = int(sys.argv[sys.argv.index("--n") + 1]) if "--n" in sys.argv else (8 if lang == "en" else 10)


TERMS = re.compile("|".join(re.escape(v) for v in sorted(VOCAB[lang], key=len, reverse=True)), re.I)


def tokens(s):
    s = TERMS.sub(" ", s)
    if lang == "en":
        return re.findall(r"[a-z0-9']+", s.lower().replace("’", "'"))
    return re.findall(r"[一-鿿A-Za-z0-9]", s.lower())


def grams(toks):
    return {tuple(toks[i:i + N]) for i in range(len(toks) - N + 1)}


items = {it["id"]: it for it in json.loads((HERE / f"items_{lang}.json").read_text(encoding="utf-8"))}
for seed in (1, 2, 3):
    run = HERE / "runs" / f"{lang}_ungated_{seed}"
    split = json.loads((run / "split.json").read_text(encoding="utf-8"))
    skill = grams(tokens((run / f"skill_r{rnd:02d}.md").read_text(encoding="utf-8")))
    hit = lambda i: bool((grams(tokens(items[i]["raw"])) | grams(tokens(items[i]["truth"]))) & skill)
    tune = [i for i in split["tune"] if hit(i)]
    unseen = [i for i in split["holdout"] + split["test"] if hit(i)]
    print(f"seed {seed}: tune {len(tune)}/{len(split['tune'])} {tune}  unseen {len(unseen)} {unseen}")