import random

import pytest

from bpe import (MIN_COUNT, NUM_MERGES, PIECE, _apply, decode, encode, load_corpus, pieces, train,
                 train_counted, vocab_bytes)
from samples import SAMPLES

CORPUS = load_corpus()
MERGES = train(CORPUS)


def test_worked_example_by_hand():
    """The lesson's 11-byte string. 'aa' is seen 4 times (positions 0-1, 1-2, 5-6, 6-7), so it is
    merge 0 (id 256). The string is now aa a b d aa a b a c; 'ab' is seen twice and the pair
    ('aa', 'a') twice too, and (97, 98) is the smaller tuple, so 'ab' is merge 1 (id 257). Then
    ('aa', 'ab') is seen twice: 'aaab' is merge 2 (id 258). Nothing is left that is seen twice."""
    learned = train_counted("aaabdaaabac")
    assert learned == [((97, 97), 4), ((97, 98), 2), ((256, 257), 2)]
    merges = [pair for pair, _ in learned]
    assert encode("aaabdaaabac", merges) == [258, 100, 258, 97, 99]
    assert vocab_bytes(merges)[256:] == [b"aa", b"ab", b"aaab"]


def test_check_yourself_example():
    """The quiz answer in the lesson. 'ababab' has 'ab' three times, so merge 0 is ab (id 256). The
    string is now 256 256 256, where (256, 256) is seen twice: merge 1 (id 257), applied left to
    right, leaves [257, 256]. The pair (257, 256) is seen once, so training stops."""
    assert train_counted("ababab") == [((97, 98), 3), ((256, 256), 2)]
    assert encode("ababab", train("ababab")) == [257, 256]


def test_tie_goes_to_the_smaller_pair_not_to_first_seen():
    """Guards a tie-break that depends on dict order: the widget and the code would disagree.
    In 'ba ab' the pairs (98, 97), (32, 97) and (97, 98) are each seen once."""
    assert train_counted("ba ab", 1, min_count=1) == [((32, 97), 1)]


def test_training_stops_when_no_pair_is_seen_min_count_times():
    assert train("abcd", 10) == []                       # every pair once
    assert len(train("abab", 10)) == 1                   # 'ab' twice, then ('ab', 'ab') once
    assert train("abcd", 10, min_count=1) != []          # the same text, a lower bar


def test_a_pair_run_is_merged_left_to_right():
    """'aaa' with the merge (97, 97) is [aa, a], never [a, aa]."""
    assert encode("aaa", [(97, 97)]) == [256, 97]


def test_pieces_cover_every_character_once():
    """Guards a gap or overlap in the six alternatives, which would drop or repeat text."""
    rng = random.Random(7)
    alphabet = list("abcXYZ 019,.;!?\t\r\n\x00\x1f\x7f") + ["é", "한", "\U0001f600", " ", "
"]
    texts = ["".join(rng.choice(alphabet) for _ in range(rng.randint(0, 40))) for _ in range(300)]
    texts += [CORPUS] + list(SAMPLES.values())
    assert all("".join(pieces(t)) == t for t in texts)
    assert any(len(pieces(t)) > 3 for t in texts)        # non-vacuous: texts do split


@pytest.mark.parametrize("text, expected", [
    ("the cat", ["the", " cat"]),
    ("a  b", ["a", " ", " b"]),                           # a run of spaces leaves one for the next word
    ("end  \n", ["end", "  \n"]),                         # trailing whitespace stays together
    ("x=12.5", ["x", "=", "12", ".", "5"]),
    ("한국어 ok", ["한국어", " ok"]),
])
def test_piece_examples(text, expected):
    assert pieces(text) == expected


def test_the_corpus_gives_all_the_merges_asked_for():
    """If the corpus shrinks below this, the lesson's '512 merges' and the widget's slider lie."""
    learned = train_counted(CORPUS)
    assert len(learned) == NUM_MERGES
    counts = [c for _, c in learned]
    assert min(counts) >= MIN_COUNT
    assert counts == sorted(counts, reverse=True)         # a merge cannot create a pair more common than itself


def test_round_trip_is_lossless_for_any_text():
    rng = random.Random(2026)
    alphabet = list("abcdefgh ,.0123456789\n\t") + ["é", "한", "국", "\U0001f600", "\x00", "\x7f"]
    texts = ["".join(rng.choice(alphabet) for _ in range(rng.randint(0, 60))) for _ in range(300)]
    texts += list(SAMPLES.values()) + [CORPUS[i:i + 400] for i in range(0, 4000, 800)]
    shorter = 0
    for t in texts:
        ids = encode(t, MERGES)
        assert decode(ids, MERGES) == t
        shorter += len(ids) < len(t.encode())
    assert shorter > 100                                  # non-vacuous: merges really fire


def test_lowest_rank_first_equals_replaying_the_merges_in_order():
    """The lesson says the two are the same thing; this is the check."""
    rng = random.Random(5)
    alphabet = list("the cat sat on a mat 0123,.\n") + ["\ud55c"]
    for _ in range(200):
        text = "".join(rng.choice(alphabet) for _ in range(rng.randint(1, 30)))
        replayed = []
        for piece in pieces(text):
            word = tuple(piece.encode("utf-8"))
            for k, pair in enumerate(MERGES):
                word = _apply(word, pair, 256 + k)
            replayed.extend(word)
        assert encode(text, MERGES) == replayed


def test_fewer_merges_never_make_text_shorter():
    sample = SAMPLES["English"] + SAMPLES["Code"]
    lengths = [len(encode(sample, MERGES[:k])) for k in (0, 64, 128, 256, 512)]
    assert lengths == sorted(lengths, reverse=True) and lengths[0] == len(sample.encode()) > lengths[-1]


def test_cut_bytes_decode_to_a_replacement_character_not_an_error():
    """A Korean syllable is 3 bytes. A token boundary can fall inside it."""
    assert decode([237, 149], []) == "�"
    assert decode([237, 149, 156], []) == "한"


def test_vocab_bytes_by_hand():
    table = vocab_bytes([(97, 98), (256, 99)])
    assert len(table) == 258 and table[97] == b"a" and table[256] == b"ab" and table[257] == b"abc"


def test_samples_are_not_in_the_corpus():
    """A sample the toy trained on would make the held-out counts look better than they are."""
    for name, text in SAMPLES.items():
        windows = {text[i:i + 30] for i in range(len(text) - 29)}
        assert windows and not any(w in CORPUS for w in windows), name


def test_piece_pattern_is_the_one_the_widget_expects():
    assert PIECE.pattern.count("|") == 5                  # six alternatives, as bpe.py and the widget say