"""Train the shared 10,000-token BPE tokenizer for en / hi / te / mr.

Design (see the write-up for the reasoning):
  * HuggingFace `tokenizers` BPE, trained on the four India pages (the evaluation
    corpus; the assignment allows any training data), each repeated `WEIGHTS[lang]` times.
  * No normaliser: NFKC would rewrite visible characters (², fi, …), which breaks the
    faithfulness rule. Text goes in exactly as written.
  * Metaspace pre-tokeniser/decoder (▁ marks a space), so spaces survive round trips and
    a token never spans two words.
  * Byte fallback: 256 <0xNN> tokens, so any character outside the vocabulary still
    round-trips exactly instead of becoming [UNK].
  * Total vocabulary = 10,000 including the byte tokens.

Usage: python train_tokenizer.py [--out tokenizer.base.json] [--weights en=1,hi=1,te=1,mr=1]
"""
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"]
VOCAB = 10_000
BYTES = [f"<0x{b:02X}>" for b in range(256)]


def corpus_iter(weights):
    # whole documents, because the evaluator encodes each file in one call: with Metaspace,
    # newlines stay inside pieces, so training must see the text the same way
    for lang in LANGS:
        text = (HERE / "corpus" / f"{lang}.faithful.txt").read_text(encoding="utf-8")
        for _ in range(weights.get(lang, 1)):
            yield text


def build(weights, vocab=VOCAB):
    tok = Tokenizer(models.BPE(byte_fallback=True))
    tok.pre_tokenizer = pre_tokenizers.Metaspace(replacement="▁", prepend_scheme="always", split=True)
    tok.decoder = decoders.Sequence([decoders.Metaspace(replacement="▁", prepend_scheme="always", split=True),
                                     decoders.ByteFallback(), decoders.Fuse()])
    trainer = trainers.BpeTrainer(vocab_size=vocab - len(BYTES), min_frequency=1, show_progress=False,
                                  special_tokens=[], limit_alphabet=100_000, max_token_length=None)
    tok.train_from_iterator(corpus_iter(weights), trainer=trainer)
    # append the byte-fallback tokens after the learned vocabulary
    spec = json.loads(tok.to_str())
    v = spec["model"]["vocab"]
    for b in BYTES:
        if b not in v:
            v[b] = len(v)
    return Tokenizer.from_str(json.dumps(spec, ensure_ascii=False))


if __name__ == "__main__":
    out = HERE / "tokenizer.base.json"
    weights = {l: 1 for l in LANGS}
    for i, a in enumerate(sys.argv):
        if a == "--out":
            out = pathlib.Path(sys.argv[i + 1])
        if a == "--weights":
            weights.update({k: int(v) for k, v in (kv.split("=") for kv in sys.argv[i + 1].split(","))})
    tok = build(weights)
    tok.save(str(out))
    print(f"saved {out.name}: vocab {tok.get_vocab_size():,}, weights {weights}")