from .train import BPETrainer
from .tokenizer import AbyssTokenizer

SAMPLE = (
    "the cat sat on the mat. the cat ran. "
    "a dog sat on the log. the dog ran fast! "
) * 50


def test_roundtrip():
    trainer = BPETrainer(vocab_size=400, verbose=False)
    merges, _ = trainer.train(SAMPLE)
    tok = AbyssTokenizer(merges, vocab_size=400)


    cases = [
        "the cat sat",
        "Hello, world!",
        "café résumé naïve",
        "emoji 🐈 and 🔥",
        "def f(x): return x + 1",
        "   leading spaces",
        "",
    ]
    for s in cases:
        assert tok.decode(tok.encode(s)) == s, f"round-trip failed: {s!r}"
    print("roundtrip: OK")


def test_specials():
    trainer = BPETrainer(vocab_size=400)
    merges, _ = trainer.train(SAMPLE)
    tok = AbyssTokenizer(merges, vocab_size=400)

    s = "<|user|>hi<|endofturn|>"
    ids = tok.encode(s)
    assert tok.special["<|user|>"] in ids
    assert tok.special["<|endofturn|>"] in ids
    assert tok.decode(ids) == s
    print("specials: OK")


def test_compression():
    trainer = BPETrainer(vocab_size=400)
    merges, _ = trainer.train(SAMPLE)
    tok = AbyssTokenizer(merges, vocab_size=400)
    text = "the cat sat on the mat"
    n_bytes = len(text.encode("utf-8"))
    n_tokens = len(tok.encode(text))
    print(f"compression: {n_bytes} bytes -> {n_tokens} tokens")
    assert n_tokens < n_bytes


if __name__ == "__main__":
    test_roundtrip()
    test_specials()
    test_compression()
    print("\nall tests passed")