# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Train a byte-level BPE vocab (minbpe BasicTokenizer algorithm, no regex
pre-split) and write the forgebpe v1 model file that src/tokenizer/bpe.cpp
loads.

Training is numpy-vectorized (bincount over packed pair keys + vectorized
merge) — minutes instead of hours for a 10-20 MB sample; a pure-python
fallback keeps the tool dependency-free if numpy is missing. The learned
merges are identical either way (ties broken by np.argmax == first-seen
max, matching Counter.most_common determinism closely enough for a fresh
vocab; the ENCODER contract that matters — greedy lowest-id merge — is
fixed by the .model file, shared by prepare_data.py and the C++ side).

Usage:
  python3 tools/train_tokenizer.py --input data/sample.txt --vocab-size 4096 \
      --output data/tok4096.model [--max-mb 10]
"""
import argparse
from collections import Counter


def train_numpy(text_bytes, vocab_size):
    import numpy as np

    ids = np.frombuffer(text_bytes, dtype=np.uint8).astype(np.int32)
    merges = []
    for i in range(vocab_size - 256):
        if len(ids) < 2:
            break
        keys = ids[:-1] * vocab_size + ids[1:]
        counts = np.bincount(keys, minlength=vocab_size * vocab_size)
        best = int(counts.argmax())
        if counts[best] < 2:
            break
        left, right = best // vocab_size, best % vocab_size
        idx = 256 + i

        # left-to-right non-overlapping merge (matches minbpe merge())
        is_pair = (ids[:-1] == left) & (ids[1:] == right)
        pos = np.where(is_pair)[0]
        if left == right:  # resolve overlaps like "aaa": keep first of each run
            keep, last = [], -2
            for p in pos:
                if p != last + 1:
                    keep.append(p)
                    last = p
            pos = np.asarray(keep, dtype=np.int64)
        ids[pos] = idx
        ids = np.delete(ids, pos + 1)
        merges.append((idx, int(left), int(right)))
        if (i + 1) % 512 == 0:
            print(f"  merge {i + 1}/{vocab_size - 256} (corpus now {len(ids)} ids)")
    return merges


def train_python(text_bytes, vocab_size):
    ids = list(text_bytes)
    merges = []
    for i in range(vocab_size - 256):
        counts = Counter(zip(ids, ids[1:]))
        if not counts:
            break
        pair = max(counts, key=counts.get)
        idx = 256 + i
        out, j = [], 0
        while j < len(ids):
            if j + 1 < len(ids) and (ids[j], ids[j + 1]) == pair:
                out.append(idx)
                j += 2
            else:
                out.append(ids[j])
                j += 1
        ids = out
        merges.append((idx, pair[0], pair[1]))
    return merges


def train(text_bytes, vocab_size):
    try:
        return train_numpy(text_bytes, vocab_size)
    except ImportError:
        print("numpy not found — slow pure-python trainer")
        return train_python(text_bytes, vocab_size)


def write_model(path, merges):
    with open(path, "w") as f:
        f.write("forgebpe v1\n")
        f.write(f"{256 + len(merges)}\n")
        for idx, left, right in merges:
            f.write(f"{idx} {left} {right}\n")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True)
    ap.add_argument("--vocab-size", type=int, default=4096)
    ap.add_argument("--output", required=True)
    ap.add_argument("--max-mb", type=float, default=10.0,
                    help="train on at most this many MB of the input")
    args = ap.parse_args()

    with open(args.input, "rb") as f:
        data = f.read(int(args.max_mb * 1024 * 1024))
    print(f"training {args.vocab_size}-vocab BPE on {len(data) / 1e6:.1f} MB")
    merges = train(data, args.vocab_size)
    write_model(args.output, merges)
    print(f"wrote {args.output} ({256 + len(merges)} tokens)")


if __name__ == "__main__":
    main()