#!/usr/bin/env python
"""Source stories -> top-k vocabulary -> filtered stories, tokenizer and training blocks on the Hub, in one run.

  export HF_TOKEN=hf_...            # needs write access
  python build_topk_dataset.py --k 300 --max-unk 8 --workers 16

One file, no local imports: the vocabulary rules of tinystories_restricted_vocab (whitelist check, top-k lexicon,
<UNK> rewrite) are copied in below. Needs only pip packages:
  pip install lemminflect ftfy datasets tokenizers "transformers>=4.45" huggingface_hub pyarrow numpy

Stages (a rerun resumes: finished story shards and the tokenizer are reused, uploads are repeated):
  1. lexicon    top-k lemmas of the first 300,000 training stories, built exactly as preprocess_corpus.py does
                (or pass --lexicon path/to/lexicon.json)
  2. stories    stream the source, ftfy, split by hash, rewrite with unk_rewrite; keep stories with status "ok"
                and at most --max-unk <UNK>s  ->  <out>/data/{train,validation}-NNNNN.parquet
  3. tokenizer  byte-level BPE on the kept training stories, train_tokenizer.py's recipe  ->  <out>/tokenizer
  4. upload     <user>/<name> (stories, private dataset) and <user>/<name>-tokenizer (private model repo)
  5. blocks     every story + <|endoftext|>, concatenated, cut into --seq-len blocks, train blocks shuffled once,
                written as --num-shards equal parquet shards  ->  <out>/blocks<seq-len>  (prepare_data.py's layout)
  6. upload     <user>/<name>-tokenized (private dataset), then prints the train.py command

train.py does not change: python train.py --data <printed hf:// path> --output_dir runs/<name>

Stories that are certain to exceed --max-unk are skipped before unk_rewrite: a lowercase word that is outside the
lexicon, has no American spelling in it, and is neither an "-ly" adverb nor an adjective always becomes one <UNK>,
so counting those words gives a lower bound. Every story that could pass goes through unk_rewrite itself, so the
kept stories are exactly what preprocess_corpus.py + prepare_data.py --max_unk would keep.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import multiprocessing as mp
import os
import re
import shutil
import sys
import time
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from functools import cached_property, lru_cache
from itertools import islice
from pathlib import Path

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")  # we parallelise with processes instead

# =====================================================================================================
# Vocabulary rules, copied from tinystories_restricted_vocab/restricted_vocab so this file runs on its own:
# vocab.py (the whitelist check), lexicon.py (top-k lexicon), mechanical.py (the <UNK> rewrite).
# =====================================================================================================

# ------------------------------------------------------------------------------------------- vocab.py
ALLOWED_SYMBOLS = set(".,;:!?'\"()[]-–—…/%$£€“”‘’")


TOKEN_RE = re.compile(
    r"(?P<ph><(?:NAME|NUM|UNK)_\d+>(?:['’]s\b)?)"
    r"|(?P<ord>\d+(?:st|nd|rd|th|s)\b)"
    r"|(?P<num>\d+(?:,\d{3})*(?:\.\d+)?)"
    r"|(?P<word>[A-Za-z]+(?:['’][A-Za-z]+)*(?:-[A-Za-z]+(?:['’][A-Za-z]+)*)*['’]?)"
    r"|(?P<sym>\S)"
)


CLITICS = ("n't", "'s", "'re", "'ve", "'ll", "'d", "'m")


def _strip_clitic(w: str) -> str:
    w = w.replace("’", "'")
    low = w.lower()
    if low == "won't":
        return w[:2] + "ill" if w[:1].isupper() else "will"
    if low in ("can't", "shan't"):
        return w[:-3]  # can't -> can, shan't -> shan (not a word)
    if w.endswith("s'"):  # plural possessive
        return w[:-1]
    for c in CLITICS:
        if w.endswith(c) and len(w) > len(c):
            return w[: -len(c)]
    return w


def _number_basic(tok: str) -> bool:
    if "," in tok:
        return False
    intpart, _, dec = tok.partition(".")
    return len(intpart) <= 4 and len(dec) <= 2


@dataclass
class Report:
    n_words: int = 0
    bad: Counter = field(default_factory=Counter)  # (kind, token) -> count
    n_placeholders: int = 0

    @property
    def n_bad(self) -> int:
        return sum(self.bad.values())

    @property
    def ok(self) -> bool:
        return self.n_bad == 0 and self.n_words > 0

    @property
    def n_bad_words(self) -> int:
        """Violations that are words/names/numbers (not stray symbols)."""
        return sum(c for (k, _), c in self.bad.items() if k != "symbol")

    @property
    def coverage(self) -> float:
        """Share of words that are allowed. Symbol violations (e.g. '*', '#')
        make the text fail `ok` but are not counted here."""
        return 1.0 - self.n_bad_words / max(1, self.n_words)

    @property
    def coverage_ignoring_names(self) -> float:
        """For ranking *source* text: capitalized unknowns would become
        placeholders anyway, so don't count them as hard words."""
        hard = sum(c for (k, _), c in self.bad.items() if k not in ("name", "number"))
        return 1.0 - hard / max(1, self.n_words)

    def bad_tokens(self, k: int = 20) -> list[str]:
        return [t for (_, t), _ in self.bad.most_common(k)]

    def as_dict(self) -> dict:
        return {
            "ok": self.ok,
            "n_words": self.n_words,
            "n_bad": self.n_bad,
            "coverage": round(self.coverage, 4),
            "bad": self.bad_tokens(50),
        }


def case_ok(part: str) -> bool:
    """Words must be lowercase, Capitalized or ALL CAPS ("learnING" is not)."""
    letters = "".join(c for c in part if c.isalpha())
    return (not letters or letters.islower() or letters.isupper()
            or (letters[0].isupper() and letters[1:].islower()))


def check(text: str, lexicon) -> Report:
    """check(text, lexicon) for the lexicon case (the only one used here), loop copied unchanged."""
    rep = Report()
    single_names, ok_fn = lexicon.names, lexicon.word_ok
    for m in TOKEN_RE.finditer(text):
        kind = m.lastgroup
        tok = m.group()
        if kind == "ph":
            rep.n_placeholders += 1
            rep.n_words += 1
        elif kind == "ord":
            rep.n_words += 1
            if not _number_basic(re.sub(r"\D", "", tok)):
                rep.bad[("number", tok)] += 1
        elif kind == "num":
            rep.n_words += 1
            if not _number_basic(tok):
                rep.bad[("number", tok)] += 1
        elif kind == "sym":
            if tok not in ALLOWED_SYMBOLS:
                rep.bad[("symbol", tok)] += 1
        else:  # word, possibly hyphenated, possibly with clitic
            for part in tok.split("-"):
                rep.n_words += 1
                if not case_ok(part):
                    rep.bad[("case", part)] += 1
                    continue
                bare = part.rstrip("'’")  # closing single quote: 'Get well, Tom'
                if (part in single_names or bare in single_names
                        or _strip_clitic(part) in single_names):
                    continue
                base = _strip_clitic(part).lower()
                if base.endswith("'"):
                    base = base[:-1]
                if not base.isalpha() or not ok_fn(base):
                    k = "name" if part[:1].isupper() else "word"
                    rep.bad[(k, part)] += 1
    return rep


# ----------------------------------------------------------------------------------------- lexicon.py
LEXICON_WORD = re.compile(r"[A-Za-z]+(?:['’][A-Za-z]+)*")


CLITIC_RE = re.compile(r"(n't|'s|'re|'ve|'ll|'d|'m)$")


POS_ORDER = ("VERB", "AUX", "NOUN", "ADJ", "ADV", "PROPN")


PRONOUN_FORMS = {
    "i": {"me", "my", "mine", "myself"},
    "you": {"your", "yours", "yourself", "yourselves"},
    "he": {"him", "his", "himself"}, "she": {"her", "hers", "herself"},
    "it": {"its", "itself"}, "we": {"us", "our", "ours", "ourselves"},
    "they": {"them", "their", "theirs", "themselves"},
}


NOT_NAMES = frozenset("""
hi hey hello wow oh ok okay yes yay bye goodbye ouch oops thanks please sorry
tv mr mrs ms dr miss sir
""".split())


@dataclass
class Lexicon:
    name: str
    lemmas: list[str]                 # top-K lemmas, most frequent first
    forms: frozenset[str]             # every allowed lowercase word form
    names: frozenset[str] = frozenset()  # allowed proper names (case-sensitive)
    top_oov: list[str] = field(default_factory=list)  # frequent words left out
    stats: dict = field(default_factory=dict)

    def word_ok(self, w: str) -> bool:
        """w is a lowercase word with clitics already stripped."""
        return w in self.forms

    @cached_property
    def prefixes(self) -> tuple[frozenset[str], frozenset[str]]:
        lp = {f[:i] for f in self.forms for i in range(1, len(f) + 1)}
        np_ = {n[:i] for n in self.names for i in range(1, len(n) + 1)}
        return frozenset(lp), frozenset(np_)

    def prompt_word_list(self) -> str:
        return ", ".join(self.lemmas)

    def save(self, path: str | Path) -> None:
        Path(path).write_text(json.dumps({
            "name": self.name, "lemmas": self.lemmas, "forms": sorted(self.forms),
            "names": sorted(self.names), "top_oov": self.top_oov, "stats": self.stats,
        }, indent=0))

    @classmethod
    def load(cls, path: str | Path) -> "Lexicon":
        d = json.loads(Path(path).read_text())
        return cls(d["name"], d["lemmas"], frozenset(d["forms"]),
                   frozenset(d["names"]), d.get("top_oov", []), d.get("stats", {}))


def _lemma_candidates(word: str) -> list[str]:
    from lemminflect import getAllLemmas

    lem = getAllLemmas(word)
    out = []
    for pos in POS_ORDER:
        for l in lem.get(pos, ()):
            if l.lower() not in out:
                out.append(l.lower())
    return out or [word]


def _inflections(lemma: str) -> set[str]:
    from lemminflect import getAllInflections

    out = {lemma}
    for forms in getAllInflections(lemma).values():
        out.update(f.lower() for f in forms)
    return out


def build_top_k(texts: list[str], k: int = 1000, n_names: int = 50,
                name: str | None = None) -> Lexicon:
    """Top-k lemmas of `texts`. A word counts as a name if it is almost never
    written in lowercase (>= 95% capitalized); the n_names most frequent names
    are allowed as-is, separately from the k lemmas."""
    # lower -> [lowercase uses, capitalized uses, capitalized mid-sentence uses]
    case = defaultdict(lambda: [0, 0, 0])
    for t in texts:
        t = t.replace("’", "'")
        for m in LEXICON_WORD.finditer(t):
            w = m.group()
            low = CLITIC_RE.sub("", w.lower()) if "'" in w else w.lower()
            if not low.isalpha():
                continue
            if w[0].islower():
                case[low][0] += 1
            else:
                case[low][1] += 1
                before = t[: m.start()].rstrip(" \t\"“”'(")
                if before and before[-1] not in ".!?:\n":
                    case[low][2] += 1
    total = {w: a + b for w, (a, b, _) in case.items()}
    # A name is (almost) never lowercase, and often capitalized mid-sentence;
    # the second test rules out "Hey", "Hi", "Once" (always sentence-initial).
    is_name = {w for w, (a, b, mid) in case.items()
               if w != "i" and w not in NOT_NAMES and a + b >= 3 and b >= 0.95 * (a + b) and mid >= 0.3 * b}

    lemma_count: Counter = Counter()
    for w, n in total.items():
        if w in is_name:
            continue
        cands = _lemma_candidates(w)
        # ambiguous forms (saw: see / saw) go to the more frequent lemma
        best = max(cands, key=lambda c: (total.get(c, 0), -cands.index(c)))
        lemma_count[best] += n

    lemmas = [l for l, _ in lemma_count.most_common(k)]
    forms = set()
    for l in lemmas:
        forms |= _inflections(l)
        forms |= PRONOUN_FORMS.get(l, set())
    names = [w.capitalize() for w, _ in sorted(((w, total[w]) for w in is_name),
                                               key=lambda x: -x[1])[:n_names]]
    all_tokens = sum(n for w, n in total.items())
    covered = sum(n for w, n in total.items() if w in forms or w.capitalize() in names)
    top_oov = [w for w, _ in sorted(((w, n) for w, n in total.items()
                                     if w not in forms and w not in is_name),
                                    key=lambda x: -x[1])[:40]]
    return Lexicon(
        name=name or f"top{k}",
        lemmas=lemmas, forms=frozenset(forms), names=frozenset(names),
        top_oov=top_oov,
        stats={"k": k, "n_texts": len(texts), "n_forms": len(forms),
               "token_coverage": round(covered / max(1, all_tokens), 4),
               "distinct_words": len(total), "n_detected_names": len(is_name)},
    )


# -------------------------------------------------------------------------------------- mechanical.py
INTENSIFIERS = {"very", "so", "really", "too", "quite", "extremely", "super"}


ADJ_SUFFIX = re.compile(r"(y|ful|ous|ish|less|ive|able|ible)$")


@lru_cache(maxsize=100_000)
def _analyze(word: str) -> tuple[str, tuple[str, ...], str | None]:
    """(lemma, POS candidates, Penn tag of this form) for a lowercase word."""
    from lemminflect import getAllInflections, getAllLemmas

    lem = getAllLemmas(word)
    pos = tuple(p for p in ("NOUN", "VERB", "ADJ", "ADV") if p in lem)
    lemma = lem[pos[0]][0].lower() if pos else word
    tag = None
    for p in pos:
        for t, forms in getAllInflections(lem[p][0]).items():
            if word in (f.lower() for f in forms):
                tag = t
                break
        if tag:
            break
    return lemma, pos, tag


_ARTICLE = re.compile(r"\b(a|an|A|An) ([A-Za-z]+)")


def fix_articles(text: str) -> str:
    """a/an agreement after words were dropped or replaced (simple vowel rule)."""
    def sub(m):
        art, nxt = m.group(1), m.group(2)
        vowel = nxt[0].lower() in "aeiou" and not nxt.lower().startswith(("uni", "use", "one"))
        new = ("an" if vowel else "a")
        if art[0].isupper():
            new = new.capitalize()
        return f"{new} {nxt}"
    return _ARTICLE.sub(sub, text)


_CAP_WORD = re.compile(r"[A-Z][a-z]+")


def source_names(source: str, lexicon=None) -> tuple[str, ...]:
    """Out-of-vocabulary names in the source: capitalized words outside the
    vocabulary that appear at least once mid-sentence (so a sentence-initial
    "Many" or "Twisty" is not mistaken for a name)."""
    rep = check(source, lexicon)
    cands = {_strip_clitic(t).rstrip("'") for (k, t) in rep.bad if k == "name"}
    names = set()
    for m in _CAP_WORD.finditer(source):
        w = m.group()
        if w in cands:
            before = source[: m.start()].rstrip(" \t\"“”'(")
            if before and before[-1] not in ".!?:\n":
                names.add(w)
    return tuple(sorted(names))


_WORD_RE = re.compile(r"[A-Za-z]+(?:['’][A-Za-z]+)*")


BRITISH = {"mummy": "mommy", "mum": "mom", "aeroplane": "airplane", "grey": "gray",
           "pyjamas": "pajamas", "tyre": "tire", "biscuit": "cookie", "jumper": "sweater"}


_SPELL = (("our", "or"), ("ise", "ize"), ("ised", "ized"), ("ising", "izing"),
          ("tre", "ter"), ("yse", "yze"), ("ll", "l"))


NON_DROPPABLE_ADVERBS = frozenset(
    "hardly barely scarcely rarely seldom nearly merely only early likely "
    "unlikely partly mostly".split())


def american_spelling(word: str, lexicon) -> str | None:
    w = word.lower()
    if w in BRITISH and BRITISH[w] in lexicon.forms:
        return BRITISH[w]
    for a, b in _SPELL:
        if a in w:
            c = w.replace(a, b)
            if c in lexicon.forms:
                return c
    return None


LINKING_VERBS = frozenset("""
am is are was were be been being feel feels felt feeling look looks looked
seem seems seemed get gets got become became stay stays stayed turn turned
""".split())


def _attributive(text: str, matches, k: int) -> bool:
    """Is matches[k] an adjective directly modifying a following noun
    ("the colourful flower"), so that dropping it keeps the sentence intact?
    Not when the next word cannot be a noun ("a camera underwater before",
    "it was cold outside") or when a linking verb comes first ("felt cold")."""
    m = matches[k]
    if text[m.end():m.end() + 1] != " " or k + 1 >= len(matches):
        return False
    nxt = matches[k + 1]
    if nxt.start() != m.end() + 1:
        return False
    nxt_low = nxt.group().lower()
    if nxt_low in ("one", "ones"):
        return False
    _, nxt_pos, nxt_tag = _analyze(_strip_clitic(nxt_low))
    if "NOUN" not in nxt_pos:
        return False
    prev = matches[k - 1].group().lower() if k > 0 else ""
    return prev not in LINKING_VERBS


def _match_case(src: str, rep: str) -> str:
    return rep[0].upper() + rep[1:] if src[:1].isupper() else rep


_MOJIBAKE = re.compile(r"â€[œ\x9d™˜”“]?|†|â€")


def clean_symbols(text: str) -> str:
    """Repair quote marks left broken after ftfy (e.g. '“', '†') and drop any
    other character the validator would reject as a stray symbol."""
    text = _MOJIBAKE.sub('"', text).replace("`", "'").replace("´", "'")
    out = []
    for ch in text:
        if ch.isalnum() or ch.isspace() or ch in ALLOWED_SYMBOLS or ch in "'<>_":
            out.append(ch)
    text = "".join(out)
    return re.sub(r"[ \t]+\n", "\n", re.sub(r"\n{3,}", "\n\n", text))


def unk_rewrite(text: str, lexicon, names: tuple[str, ...] = (),
                synonyms: dict[str, str] | None = None,
                name_map: dict[str, str] | None = None) -> tuple[str, Counter]:
    """Bring a story into the lexicon without generating text:

      names outside the vocabulary -> name_map[name] if given, else <NAME_k>
      British spellings            -> American ones in the vocabulary
      synonyms[lemma] (optional, a checked dictionary) -> that word, inflected
      adverbs                      -> dropped
      adjectives before a word     -> dropped (with an intensifier in front)
      everything else              -> <UNK_k>, one number per lemma per story

    Returns (new text, Counter of what happened to each out-of-vocabulary word;
    keys starting with "@" count the lemmas that became UNK)."""
    text = clean_symbols(text)
    names = tuple(names) or source_names(text, lexicon)
    synonyms = synonyms or {}
    stats: Counter = Counter()
    pieces, last = [], 0
    name_ids: dict[str, int] = {}
    unk_ids: dict[str, int] = {}
    matches = list(_WORD_RE.finditer(text))
    drop_next_space = False
    for k, m in enumerate(matches):
        w = m.group()
        base = _strip_clitic(w).replace("’", "'")
        clitic = w[len(base):]
        if check(w, lexicon).ok:
            continue
        rep = None
        if base in names:
            stats["name"] += 1
            if name_map and base in name_map:
                rep = name_map[base]
            else:
                name_ids.setdefault(base, len(name_ids) + 1)
                rep = f"<NAME_{name_ids[base]}>"
        else:
            low = base.lower()
            am = american_spelling(low, lexicon)
            lemma, pos, tag = _analyze(low)
            nxt = text[m.end():m.end() + 2]
            # Only -ly manner adverbs are safe to drop ("ran quickly"); others can
            # carry the meaning ("it got him nowhere", "she hardly ate").
            is_adv = (((pos and set(pos) <= {"ADV"}) or not pos) and low.endswith("ly")
                      and low not in NON_DROPPABLE_ADVERBS)
            is_adj = ("ADJ" in pos and "VERB" not in pos) or (not pos and ADJ_SUFFIX.search(low))
            if am:
                stats["spelling"] += 1
                rep = _match_case(base, am)
            elif lemma in synonyms:
                syn = _inflect(synonyms[lemma], tag) or synonyms[lemma]
                stats["synonym"] += 1
                rep = _match_case(base, syn)
            elif is_adv or (is_adj and _attributive(text, matches, k)):
                stats["drop_adverb" if is_adv else "drop_adjective"] += 1
                # cut the word and the space after it (or before, at a clause end)
                start, end = m.start(), m.end()
                if text[end:end + 1] == " ":
                    end += 1
                elif start > 0 and text[start - 1] == " ":
                    start -= 1
                if is_adj and k > 0 and matches[k - 1].group().lower() in INTENSIFIERS \
                        and text[matches[k - 1].end():m.start()] == " ":
                    start = matches[k - 1].start()
                pieces.append(text[last:start])
                last = end
                continue
            else:
                kind = ("unk_predicate_adjective" if is_adj
                        else "unk_noun" if (tag or "").startswith("NN")
                        else "unk_verb" if (tag or "").startswith("VB")
                        else "unk_other")
                stats[kind] += 1
                stats["@" + lemma] += 1  # which lemmas became UNK ("@" keys)
                unk_ids.setdefault(lemma, len(unk_ids) + 1)
                rep = f"<UNK_{unk_ids[lemma]}>"
        pieces.append(text[last:m.start()])
        pieces.append(rep + clitic)
        last = m.end()
    pieces.append(text[last:])
    out = fix_articles("".join(pieces))
    out = re.sub(r"(?<=\w) ([,.!?])", r"\1", out)  # "the dog ." after a drop
    return out, stats


def _inflect(lemma: str, tag: str | None) -> str | None:
    """Repairer._inflect from mechanical.py (only reached when a synonyms dictionary is passed, which this script
    never does)."""
    from lemminflect import getInflection

    if not tag:
        return lemma
    forms = getInflection(lemma, tag=tag)
    return forms[0] if forms else None


# ------------------------------------------------------------------ helpers copied from preprocess_corpus.py
SCHEMA_FIELDS = [
    ("id", "int64"), ("text", "string"), ("original", "string"), ("n_words", "int32"),
    ("n_oov", "int32"), ("n_unk", "int32"), ("n_unk_distinct", "int32"),
    ("n_names", "int32"), ("n_dropped", "int32"), ("unk_lemmas", "list<string>"),
    ("status", "string"),
]


def _schema(keep_original: bool):
    import pyarrow as pa

    types = {"int64": pa.int64(), "int32": pa.int32(), "string": pa.string(),
             "list<string>": pa.list_(pa.string())}
    return pa.schema([(n, types[t]) for n, t in SCHEMA_FIELDS if keep_original or n != "original"])


def write_parquet(rows: list[dict], path: Path, schema) -> None:
    import pyarrow as pa
    import pyarrow.parquet as pq

    cols = {f.name: [r[f.name] for r in rows] for f in schema}
    table = pa.Table.from_pydict(cols, schema=schema)
    tmp = path.with_suffix(".parquet.tmp")
    pq.write_table(table, tmp, compression="zstd")
    os.replace(tmp, path)


def shard_paths(out: Path, k: int) -> tuple[Path, Path]:
    return (out / "data" / f"train-{k:05d}.parquet", out / "data" / f"validation-{k:05d}.parquet")


def shard_done(out: Path, k: int) -> bool:
    return all(p.exists() for p in shard_paths(out, k))


def split_of(text: str, val_frac: float) -> str:
    h = int(hashlib.md5(text.encode("utf-8")).hexdigest()[:8], 16)
    return "validation" if h % 1_000_000 < val_frac * 1_000_000 else "train"


def iter_texts(args):
    """Source stories in a fixed order, streamed from the Hub."""
    from datasets import load_dataset

    for row in load_dataset(args.source, split="train", streaming=True):
        yield row["text"]


def get_lexicon(args) -> Path:
    """<out>/lexicon.json: --lexicon if given, else the top-k lemmas of the first --lexicon-sample training stories."""
    import ftfy

    path = args.out / "lexicon.json"
    if args.lexicon:
        Lexicon.load(args.lexicon).save(path)
        print(f"using lexicon {args.lexicon}")
        return path
    if path.exists():
        print(f"reusing {path}")
        return path
    print(f"building top-{args.k} lexicon from {args.lexicon_sample:,} training stories...", flush=True)
    texts = []
    for t in iter_texts(args):
        t = ftfy.fix_text(t or "")
        if split_of(t.strip(), args.val_frac) == "train":
            texts.append(t)
            if len(texts) >= args.lexicon_sample:
                break
    lex = build_top_k(texts, args.k, args.n_names, name=f"{args.source.split('/')[-1]}_top{args.k}")
    lex.save(path)
    print(f"token coverage of the new lexicon: {lex.stats['token_coverage']:.1%}")
    return path


# -------------------------------------------------------------------- helpers copied from train_tokenizer.py
EOT = "<|endoftext|>"


def placeholder_tokens(n_unk: int, n_name: int, n_num: int) -> list[str]:
    """Each placeholder twice: with a leading space and without, so it is always one token."""
    base = ([f"<UNK_{i}>" for i in range(1, n_unk + 1)]
            + [f"<NAME_{i}>" for i in range(1, n_name + 1)]
            + [f"<NUM_{i}>" for i in range(1, n_num + 1)])
    return [" " + b for b in base] + base


def train_bpe(texts: list[str], vocab_size: int, specials: list[str]):
    from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers

    tok = Tokenizer(models.BPE())
    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tok.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=[EOT] + specials,
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        min_frequency=2,
        show_progress=False,
    )
    tok.train_from_iterator(texts, trainer=trainer)
    return tok


def evaluate_tokenizer(tok, texts: list[str], lexicon_forms: list[str] | None) -> dict:
    enc = tok.encode_batch(texts)
    n_tok = sum(len(e.ids) for e in enc)
    n_words = sum(len(t.split()) for t in texts)
    roundtrip = sum(tok.decode(e.ids, skip_special_tokens=False) == t for e, t in zip(enc, texts))
    out = {
        "vocab_size": tok.get_vocab_size(),
        "tokens_per_story": round(n_tok / len(texts), 1),
        "tokens_per_word": round(n_tok / n_words, 3),
        "roundtrip_exact": f"{roundtrip}/{len(texts)}",
    }
    if lexicon_forms:
        single = sum(len(tok.encode(" " + f).ids) == 1 for f in lexicon_forms)
        out["lexicon_forms_single_token"] = round(single / len(lexicon_forms), 4)
    return out

UNK_RE = re.compile(r"<UNK_(\d+)>")
WORD = re.compile(r"[A-Za-z]+(?:['’][A-Za-z]+)*")  # the word pattern unk_rewrite walks over
_LEX, _MAX_UNK, _VAL_FRAC, _CACHE = None, 8, 0.005, {}


def parse_args():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--source", default="karpathy/tinystories-gpt4-clean", help="Hub dataset with the source stories")
    ap.add_argument("--k", type=int, default=300, help="number of lemmas in the vocabulary")
    ap.add_argument("--max-unk", type=int, default=8, help="keep stories with at most this many <UNK>s")
    ap.add_argument("--lexicon", default=None, help="use this lexicon.json instead of building the top-k one")
    ap.add_argument("--n-names", type=int, default=50)
    ap.add_argument("--lexicon-sample", type=int, default=300_000, help="training stories used to build the lexicon")
    ap.add_argument("--val-frac", type=float, default=0.005)
    ap.add_argument("--vocab-size", type=int, default=4096, help="BPE target; training stops earlier if it runs "
                                                                   "out of merges (about 2,700 at k=300)")
    ap.add_argument("--seq-len", type=int, default=512)
    ap.add_argument("--num-shards", type=int, default=16,
                    help="train shards; a multiple of GPUs x dataloader workers (16 suits 1-4 GPUs x 4 workers)")
    ap.add_argument("--shuffle-seed", type=int, default=0)
    ap.add_argument("--name", default=None, help="repo base name (default: tinystories-top<k>)")
    ap.add_argument("--user", default=None, help="Hub user or organization (default: the token's user)")
    ap.add_argument("--out", default=None, help="local working folder (default: out/<name>)")
    ap.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 2) - 1))
    ap.add_argument("--shard-size", type=int, default=100_000, help="source stories per story shard")
    ap.add_argument("--batch-size", type=int, default=500)
    ap.add_argument("--limit", type=int, default=0, help="stop after this many source stories (quick tests)")
    ap.add_argument("--no-upload", action="store_true", help="build everything locally, upload nothing")
    ap.add_argument("--public", action="store_true", help="create the repos as public instead of private")
    args = ap.parse_args()
    args.name = args.name or f"tinystories-top{args.k}"
    args.out = Path(args.out or f"out/{args.name}")
    return args


# ---------------------------------------------------------------------------------------------- stories
def _init_worker(lexicon_path: str, max_unk: int, val_frac: float):
    global _LEX, _MAX_UNK, _VAL_FRAC

    _LEX, _MAX_UNK, _VAL_FRAC = Lexicon.load(lexicon_path), max_unk, val_frac


def _word_status(w: str) -> int:
    """0 allowed; 1 might stay out of <UNK> (name, respelling, dropped adverb/adjective); 2 always one <UNK>."""
    s = _CACHE.get(w)
    if s is None:
        if check(w, _LEX).ok:
            s = 0
        elif w[:1].isupper():
            s = 1
        else:
            low = _strip_clitic(w).replace("’", "'").lower()
            _, pos, _ = _analyze(low)
            is_adv = (((pos and set(pos) <= {"ADV"}) or not pos) and low.endswith("ly")
                      and low not in NON_DROPPABLE_ADVERBS)
            is_adj = ("ADJ" in pos and "VERB" not in pos) or (not pos and ADJ_SUFFIX.search(low))
            s = 1 if (american_spelling(low, _LEX) or is_adv or is_adj) else 2
        _CACHE[w] = s
    return s


def _process_batch(batch: list[tuple[int, str]]) -> tuple[list[dict], Counter]:
    import ftfy

    out, seen = [], Counter()
    for idx, raw in batch:
        text = ftfy.fix_text(raw or "").strip()  # same normalisation as preprocess_corpus.py
        split = split_of(text, _VAL_FRAC)
        seen[split] += 1
        if not re.search(r"[A-Za-z]", text):
            continue
        certain = 0
        for w in WORD.findall(text):
            if _word_status(w) == 2:
                certain += 1
                if certain > _MAX_UNK:
                    break
        if certain > _MAX_UNK:
            continue
        new, st = unk_rewrite(text, _LEX)
        counts = {k: v for k, v in st.items() if not k.startswith("@")}
        n_unk = sum(v for k, v in counts.items() if k.startswith("unk_"))
        if n_unk > _MAX_UNK:
            continue
        if not check(new, _LEX).ok:  # status "invalid" in preprocess_corpus.py; prepare_data.py drops these
            seen["invalid"] += 1
            continue
        seen[f"kept_{split}"] += 1
        out.append({"id": idx, "text": new, "original": text, "n_words": len(text.split()),
                    "n_oov": sum(counts.values()), "n_unk": n_unk, "n_unk_distinct": len(set(UNK_RE.findall(new))),
                    "n_names": counts.get("name", 0),
                    "n_dropped": counts.get("drop_adverb", 0) + counts.get("drop_adjective", 0),
                    "unk_lemmas": [k[1:] for k in st if k.startswith("@")], "status": "ok", "_split": split})
    return out, seen


def build_stories(args, lex_path: Path) -> Counter:
    out = args.out
    (out / "data").mkdir(parents=True, exist_ok=True)
    schema = _schema(True)
    source = iter_texts(args)
    texts = islice(source, args.limit) if args.limit else source
    totals, t0, k = Counter(), time.time(), 0
    with mp.get_context("spawn").Pool(args.workers, initializer=_init_worker,
                                      initargs=(str(lex_path), args.max_unk, args.val_frac)) as pool:
        while True:
            block = list(islice(texts, args.shard_size))
            if not block:
                break
            stats_path = out / "data" / f".stats-{k:05d}.json"
            tr_path, va_path = shard_paths(out, k)
            if shard_done(out, k) and stats_path.exists():
                totals.update(json.loads(stats_path.read_text()))
                k += 1
                continue
            first = k * args.shard_size
            batches = [[(first + i + j, t) for j, t in enumerate(block[i:i + args.batch_size])]
                       for i in range(0, len(block), args.batch_size)]
            rows, seen = [], Counter()
            for r, s in pool.imap(_process_batch, batches):  # ordered, so rows stay in source order
                rows += r
                seen.update(s)
            write_parquet([r for r in rows if r["_split"] == "validation"], va_path, schema)
            write_parquet([r for r in rows if r["_split"] == "train"], tr_path, schema)
            stats_path.write_text(json.dumps(seen))
            totals.update(seen)
            print(f"story shard {k:05d}: kept {seen['kept_train']:,} train + {seen['kept_validation']:,} validation "
                  f"of {len(block):,} | {sum(totals[s] for s in ('train', 'validation')) / (time.time() - t0):,.0f} "
                  f"stories/s", flush=True)
            k += 1
    source.close()
    return totals


def read_texts(out: Path, split: str) -> list[str]:
    import pyarrow.parquet as pq

    return [x for f in sorted((out / "data").glob(f"{split}-*.parquet"))
            for x in pq.read_table(f, columns=["text"]).column("text").to_pylist()]


def write_story_card(args, totals: Counter, lexicon: dict) -> dict:
    stats = {s: {"source_stories": totals[s], "kept": totals[f"kept_{s}"]} for s in ("train", "validation")}
    stats.update(k=args.k, max_unk=args.max_unk, source=args.source, invalid_dropped=totals["invalid"],
                 lexicon={"lemmas": len(lexicon["lemmas"]), "forms": len(lexicon["forms"]),
                          "names": len(lexicon["names"])})
    (args.out / "stats.json").write_text(json.dumps(stats, indent=1))
    tr, va = stats["train"], stats["validation"]
    (args.out / "README.md").write_text(f"""---
configs:
- config_name: default
  data_files:
  - split: train
    path: data/train-*.parquet
  - split: validation
    path: data/validation-*.parquet
---

# {args.name}: stories in a {args.k}-lemma vocabulary, at most {args.max_unk} `<UNK>` each

Source: `{args.source}`, rewritten with `unk_rewrite` (restricted_vocab/mechanical.py, copied into build_topk_dataset.py) against `lexicon.json`
({stats['lexicon']['lemmas']} lemmas, {stats['lexicon']['forms']} word forms, {stats['lexicon']['names']} names). Kept: stories
whose rewrite passes the vocabulary check and has at most {args.max_unk} `<UNK>`s.

Train: {tr['kept']:,} of {tr['source_stories']:,} stories. Validation (about {args.val_frac:.1%} of stories, chosen by a
hash of the text): {va['kept']:,} of {va['source_stories']:,}.
""")
    return stats


# -------------------------------------------------------------------------------------------- tokenizer
def build_tokenizer(args):
    from transformers import AutoTokenizer, PreTrainedTokenizerFast

    tok_dir = args.out / "tokenizer"
    if (tok_dir / "tokenizer.json").exists():
        print(f"reusing {tok_dir}")
        return AutoTokenizer.from_pretrained(tok_dir)
    texts, val = read_texts(args.out, "train"), read_texts(args.out, "validation")[:2000]
    if not val:  # tiny test runs: hold out some training stories, as train_tokenizer.py does
        import random

        random.Random(0).shuffle(texts)
        val, texts = texts[:1000], texts[1000:]
    forms = json.loads((args.out / "lexicon.json").read_text())["forms"]
    specials = placeholder_tokens(64, 16, 8)
    over = [m for t in texts for m in re.findall(r"<(UNK|NAME|NUM)_(\d+)>", t)
            if int(m[1]) > {"UNK": 64, "NAME": 16, "NUM": 8}[m[0]]]
    if over:
        print(f"note: {len(over)} placeholders exceed the reserved range and will be split")
    print(f"training the tokenizer on {len(texts):,} stories...", flush=True)
    bpe = train_bpe(texts, args.vocab_size, specials)
    stats = evaluate_tokenizer(bpe, val, forms)
    print(json.dumps(stats))
    tok = PreTrainedTokenizerFast(tokenizer_object=bpe, bos_token=EOT, eos_token=EOT, pad_token=EOT,
                                  additional_special_tokens=specials, model_max_length=1024,
                                  clean_up_tokenization_spaces=False)
    tok.save_pretrained(tok_dir)
    (tok_dir / "tokenizer_stats.json").write_text(json.dumps(stats, indent=1))
    shutil.copy(args.out / "lexicon.json", tok_dir / "lexicon.json")
    sample = "<NAME_1> saw the <UNK_1> and the <UNK_12>. <|endoftext|>"
    ids = tok(sample)["input_ids"]
    assert tok.decode(ids) == sample, tok.decode(ids)
    assert tok.convert_tokens_to_ids(" <UNK_12>") in ids
    return tok


# ----------------------------------------------------------------------------------------------- blocks
def write_shards(tokens, seq_len, out_dir: Path, split, num_shards, shuffle_seed=None) -> dict:
    """Cut the token stream into seq_len blocks and write num_shards equal parquet files (as prepare_data.py)."""
    import numpy as np
    import pyarrow as pa
    import pyarrow.parquet as pq

    n_blocks = tokens.size // seq_len
    per_shard = n_blocks // num_shards
    if per_shard == 0:
        raise SystemExit(f"{split}: only {n_blocks} blocks, fewer than --num-shards {num_shards}")
    blocks = tokens[: n_blocks * seq_len].reshape(n_blocks, seq_len)
    order = (np.random.default_rng(shuffle_seed).permutation(n_blocks) if shuffle_seed is not None
             else np.arange(n_blocks))
    files = []
    for i in range(num_shards):
        rows = blocks[order[i * per_shard:(i + 1) * per_shard]]
        column = pa.FixedSizeListArray.from_arrays(pa.array(rows.reshape(-1)), seq_len)
        name = f"{split}-{i:05d}-of-{num_shards:05d}.parquet"
        pq.write_table(pa.table({"input_ids": column}), out_dir / name, compression="zstd", row_group_size=1024)
        files.append(name)
    return {"blocks": per_shard * num_shards, "blocks_dropped": n_blocks - per_shard * num_shards,
            "num_shards": num_shards, "blocks_per_shard": per_shard, "files": files}


def build_blocks(args, tok, totals: Counter, repos: dict, revisions: dict) -> Path:
    import numpy as np

    name = f"blocks{args.seq_len}"
    bdir = args.out / name
    meta_path = bdir / "meta.json"
    if meta_path.exists():
        old = json.loads(meta_path.read_text())
        if old.get("dataset_revision") == revisions["stories"] and old.get("tokenizer_revision") == revisions["tokenizer"]:
            print(f"reusing {bdir}")
            return bdir
    if len(tok) > np.iinfo(np.uint16).max + 1:
        raise SystemExit(f"vocab of {len(tok)} does not fit in uint16")
    shutil.rmtree(bdir, ignore_errors=True)
    bdir.mkdir(parents=True)
    tok.save_pretrained(bdir / "tokenizer")
    meta = {
        "format": "packed-blocks-v1", "seq_len": args.seq_len, "column": "input_ids", "dtype": "uint16",
        "dataset": repos["stories"], "dataset_revision": revisions["stories"],
        "tokenizer": repos["tokenizer"], "tokenizer_revision": revisions["tokenizer"], "text_column": "text",
        "filter": {"status_ok_only": True, "max_unk": args.max_unk}, "max_stories": None, "vocab_size": len(tok),
        "eos_token_id": tok.eos_token_id, "bos_token_id": tok.bos_token_id, "shuffle_seed": args.shuffle_seed,
        "layout": "stories concatenated, <|endoftext|> after each story, cut into seq_len blocks", "splits": {},
        "k": args.k, "source": args.source,
    }
    for split, shards, seed in [("train", args.num_shards, args.shuffle_seed), ("validation", 1, None)]:
        texts = read_texts(args.out, split)
        parts = []
        for i in range(0, len(texts), 20_000):
            for ids in tok(texts[i:i + 20_000], add_special_tokens=False)["input_ids"]:
                parts.append(np.asarray(ids + [tok.eos_token_id], dtype=np.uint16))
        tokens = np.concatenate(parts)
        info = write_shards(tokens, args.seq_len, bdir, split, shards, seed)
        meta["splits"][split] = {"stories_raw": totals[split], "stories": len(texts), "tokens": int(tokens.size), **info}
        print(f"{split:10s} {len(texts):,} stories -> {tokens.size:,} tokens -> {info['blocks']:,} blocks of "
              f"{args.seq_len} in {shards} shard(s)", flush=True)
    meta["created"] = time.strftime("%Y-%m-%d %H:%M:%S")
    meta_path.write_text(json.dumps(meta, indent=2))
    return bdir


def blocks_card(meta: dict, name: str) -> str:
    tr, va = meta["splits"]["train"], meta["splits"]["validation"]
    return f"""---
configs:
- config_name: {name}
  data_files:
  - split: train
    path: {name}/train-*.parquet
  - split: validation
    path: {name}/validation-*.parquet
---

# Tokenized {meta['dataset'].split('/')[-1]}: {meta['seq_len']}-token training blocks (`{name}/`)

- Source stories: `{meta['dataset']}@{meta['dataset_revision']}` (top-{meta['k']} lemmas, at most
  {meta['filter']['max_unk']} `<UNK>` per story)
- Tokenizer: `{meta['tokenizer']}@{meta['tokenizer_revision']}` (vocab {meta['vocab_size']}; a copy is in `{name}/tokenizer/`)
- Layout: every story + `<|endoftext|>` (id {meta['eos_token_id']}), concatenated, cut into {meta['seq_len']}-token
  blocks, one per row in column `input_ids` (uint16). Train blocks are shuffled once (seed {meta['shuffle_seed']}).
- train: {tr['stories']:,} stories, {tr['blocks']:,} blocks in {tr['num_shards']} shards; validation: {va['blocks']:,} blocks

Train with `train.py --data hf://datasets/<this repo>@<commit>/{name}`.
"""


# ----------------------------------------------------------------------------------------------- upload
def check_write_access(api) -> str:
    try:
        who = api.whoami()
    except Exception as e:  # noqa: BLE001 - any auth failure means no upload
        raise SystemExit(f"uploading needs a Hugging Face token with write access (export HF_TOKEN=...): {e}")
    if who.get("auth", {}).get("accessToken", {}).get("role") == "read":
        raise SystemExit("HF_TOKEN is read-only; uploading needs a write (or write-scoped fine-grained) token. "
                         "Use --no-upload to build everything locally.")
    return who["name"]


def upload_folder(api, repo, repo_type, folder, private, message, **kw) -> str:
    api.create_repo(repo, repo_type=repo_type, private=private, exist_ok=True)
    commit = api.upload_folder(repo_id=repo, repo_type=repo_type, folder_path=str(folder), commit_message=message, **kw)
    return commit.oid


# ------------------------------------------------------------------------------------------------- main
def main():
    args = parse_args()
    api = None
    user = args.user
    if not args.no_upload:
        from huggingface_hub import HfApi

        api = HfApi()
        me = check_write_access(api)  # fail before the slow part if the token cannot upload
        user = user or me
        print(f"will upload to {user}/ as {me}")
    user = user or "local"
    repos = {"stories": f"{user}/{args.name}", "tokenizer": f"{user}/{args.name}-tokenizer",
             "tokenized": f"{user}/{args.name}-tokenized"}
    args.out.mkdir(parents=True, exist_ok=True)

    t0 = time.time()
    lex_path = get_lexicon(args)  # builds (or reuses) <out>/lexicon.json
    lexicon = json.loads(Path(lex_path).read_text())
    print(f"lexicon: {len(lexicon['lemmas'])} lemmas, {len(lexicon['forms'])} forms, {len(lexicon['names'])} names")

    totals = build_stories(args, Path(lex_path))
    stats = write_story_card(args, totals, lexicon)
    print(f"stories: kept {stats['train']['kept']:,} of {stats['train']['source_stories']:,} train and "
          f"{stats['validation']['kept']:,} of {stats['validation']['source_stories']:,} validation "
          f"({time.time() - t0:.0f}s)")

    tok = build_tokenizer(args)
    revisions = {"stories": "local", "tokenizer": "local"}
    if api:
        revisions["stories"] = upload_folder(api, repos["stories"], "dataset", args.out, not args.public,
                                             f"Top-{args.k} stories with at most {args.max_unk} <UNK>",
                                             allow_patterns=["data/*.parquet", "lexicon.json", "stats.json", "README.md"])
        revisions["tokenizer"] = upload_folder(api, repos["tokenizer"], "model", args.out / "tokenizer",
                                               not args.public, f"Tokenizer for {args.name}")
        print(f"uploaded https://huggingface.co/datasets/{repos['stories']} and https://huggingface.co/{repos['tokenizer']}")

    bdir = build_blocks(args, tok, totals, repos, revisions)
    meta = json.loads((bdir / "meta.json").read_text())
    (args.out / f"{bdir.name}_README.md").write_text(blocks_card(meta, bdir.name))
    steps = round(20 * meta["splits"]["train"]["blocks"] / 1280)
    if api:
        api.create_repo(repos["tokenized"], repo_type="dataset", private=not args.public, exist_ok=True)
        api.upload_folder(repo_id=repos["tokenized"], repo_type="dataset", folder_path=str(bdir),
                          path_in_repo=bdir.name, delete_patterns="*",
                          commit_message=f"Upload {bdir.name}: {meta['splits']['train']['blocks']:,} train blocks")
        sha = api.upload_file(path_or_fileobj=blocks_card(meta, bdir.name).encode(), path_in_repo="README.md",
                              repo_id=repos["tokenized"], repo_type="dataset", commit_message=f"Describe {bdir.name}").oid
        data = f"hf://datasets/{repos['tokenized']}@{sha}/{bdir.name}"
    else:
        data = str(bdir)
    print(f"\ndone in {(time.time() - t0) / 60:.1f} min. Train with the unchanged train.py "
          f"(~{steps:,} steps at the default 20 epochs):\n  python train.py --data {data} --output_dir runs/{args.name}")


if __name__ == "__main__":
    # Streaming readers can leave threads alive at exit (see preprocess_corpus.py); exit explicitly.
    code = 0
    try:
        main()
    except SystemExit as e:
        if e.code not in (None, 0):
            print(e.code if isinstance(e.code, str) else f"exit {e.code}", file=sys.stderr)
            code = e.code if isinstance(e.code, int) else 1
    except BaseException:
        import traceback

        traceback.print_exc()
        code = 1
    sys.stdout.flush()
    sys.stderr.flush()
    os._exit(code)