"""Bring the vocabulary back to exactly 10,000 without changing the score.

Equalising removes merges, so the vocabulary shrinks. The assignment asks for a
10,000-token vocabulary, so we append merges learned from held-out Wikipedia
articles (never the India pages), keeping only merges that cannot fire on the
four evaluation files:

  * both parts are already tokens, and the merged token is new;
  * the pair never occurs side by side in the final tokenisation of any of the
    four files.

An appended merge has the lowest priority, so it can only act on the final token
sequence; if its pair never appears there, every evaluation count is unchanged.
This is re-measured, not assumed. On other text the extra merges help, since
they are common words from held-out articles.

Usage: python fill_vocab.py tokenizer.eq.json tokenizer.json
"""
import json
import pathlib
import sys

from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers

HERE = pathlib.Path(__file__).resolve().parent
LANGS = ["en", "hi", "te", "mai"]
TARGET = 10_000


def heldout_merges():
    """merges learned (in order) from the held-out articles: other Indian topics, en/hi/te/mr"""
    tok = Tokenizer(models.BPE(byte_fallback=True))
    tok.pre_tokenizer = pre_tokenizers.Metaspace(replacement="▁", prepend_scheme="always", split=True)
    tr = trainers.BpeTrainer(vocab_size=30_000, min_frequency=3, show_progress=False, limit_alphabet=2000)
    files = sorted((HERE / "heldout").glob("*__*.md"))
    tok.train_from_iterator((f.read_text(encoding="utf-8") for f in files), trainer=tr)
    return [tuple(m) if isinstance(m, list) else tuple(m.split(" ")) for m in json.loads(tok.to_str())["model"]["merges"]]


def counts(tok, texts):
    return {l: len(tok.encode(t).ids) for l, t in texts.items()}


if __name__ == "__main__":
    src, dst = sys.argv[1], sys.argv[2]
    texts = {l: (HERE / "corpus" / f"{l}.faithful.txt").read_text(encoding="utf-8") for l in LANGS}
    spec = json.loads(pathlib.Path(src).read_text(encoding="utf-8"))
    tok = Tokenizer.from_str(json.dumps(spec, ensure_ascii=False))
    before = counts(tok, texts)

    # every adjacent pair in the final tokenisation of the four files (a superset of what matters)
    present = set()
    for t in texts.values():
        toks = tok.encode(t).tokens
        present.update(zip(toks, toks[1:]))

    m = spec["model"]
    vocab, merges = m["vocab"], m["merges"]
    added = 0
    for a, b in heldout_merges():
        if len(vocab) >= TARGET:
            break
        if a in vocab and b in vocab and (a + b) not in vocab and (a, b) not in present:
            vocab[a + b] = len(vocab)
            merges.append([a, b])
            added += 1
    # contiguous ids: learned tokens, then the new ones (byte tokens keep their place)
    m["vocab"] = {t: i for i, (t, _) in enumerate(sorted(vocab.items(), key=lambda kv: kv[1]))}
    out = Tokenizer.from_str(json.dumps(spec, ensure_ascii=False))
    after = counts(out, texts)
    assert after == before, f"evaluation counts changed: {before} → {after}"
    assert out.get_vocab_size() == TARGET, out.get_vocab_size()
    pathlib.Path(dst).write_text(json.dumps(spec, ensure_ascii=False, indent=1), encoding="utf-8")
    print(f"added {added} held-out merges → vocab {out.get_vocab_size():,}; evaluation counts unchanged: {after}")