import random
import re
import string

import pytest

from bpe import NUM_MERGES, load_corpus, pieces
from letters import (DEFAULT_LETTER, DEFAULT_WORD, EXAMPLES, MAX_WORD, SEPARATOR, chunks, count_from_chunks,
                     count_ids_that_are_the_letter, count_inside, share_in_longer_chunks, spelled_out, toy_merges)

MERGES = toy_merges()
CORPUS = load_corpus()

# The lesson's hand example: three merges, in the order they were learned. er is id 256, ry is 257, st is 258.
TINY = [(101, 114), (114, 121), (115, 116)]


def test_worked_example_by_hand():
    """The Go-deeper walk-through. e+r is the first merge, so it fires first; r+y is next (the r of 'er' is
    already gone, so only the last r can join the y); s+t is last. The ids end as [258, 114, 97, 119, 98, 256, 257].
    Counting ids equal to 114 finds one r; adding up the spellings finds three."""
    rows = [[115, 116, 114, 97, 119, 98, 101, 114, 114, 121], [115, 116, 114, 97, 119, 98, 256, 114, 121],
            [115, 116, 114, 97, 119, 98, 256, 257], [258, 114, 97, 119, 98, 256, 257]]
    assert [[i for i, _ in chunks("strawberry", TINY[:k])] for k in range(4)] == rows
    assert [i for i, _ in chunks("strawberry", TINY)] == [258, 114, 97, 119, 98, 256, 257]
    assert [s for _, s in chunks("strawberry", TINY)] == ["st", "r", "a", "w", "b", "er", "ry"]
    assert count_ids_that_are_the_letter("strawberry", "r", TINY) == 1
    assert count_from_chunks("strawberry", "r", TINY) == 3


def test_check_yourself_berry_example():
    """Quiz 3: with the tiny table, berry is [98, 256, 257] (merge 2, s+t, does not apply); looking for 114 finds
    none, and the chunks er and ry hold two r's."""
    assert [i for i, _ in chunks("berry", TINY)] == [98, 256, 257]
    assert count_ids_that_are_the_letter("berry", "r", TINY) == 0
    assert count_from_chunks("berry", "r", TINY) == 2


def test_lesson_numbers():
    """Guards the figures the lesson quotes. A change to bpe.py, corpus.txt or the word list moves them, and the
    lesson text, figures and widget would then disagree with this code."""
    assert (DEFAULT_WORD, DEFAULT_LETTER) == ("strawberry", "r")
    assert len(MERGES) == NUM_MERGES == 512
    got = chunks("strawberry", MERGES)
    assert [i for i, _ in got] == [292, 114, 97, 119, 98, 263, 587]
    assert [s for _, s in got] == ["st", "r", "a", "w", "b", "er", "ry"]
    assert count_from_chunks("strawberry", "r", MERGES) == 3
    assert count_ids_that_are_the_letter("strawberry", "r", MERGES) == 1
    words = re.findall(r"[A-Za-z]+", CORPUS)
    assert len(words) == 2231 and sum(map(len, words)) == 9227
    assert round(100 * share_in_longer_chunks(words, MERGES), 1) == 72.7


@pytest.mark.parametrize("word,letter", EXAMPLES.items())
def test_every_widget_example_counts_the_same_both_ways(word, letter):
    """Guards an example whose chunk spellings and plain str.count disagree (a chunk dropped or split wrong)."""
    assert count_from_chunks(word, letter, MERGES) == word.lower().count(letter.lower())


def test_chunks_partition_the_text_for_random_ascii():
    """Guards a tokenizer that loses or repeats a character: the spellings of the chunks, joined, must be the text
    (spaces, hyphens and capitals included), and the letter count from chunks must equal str.count."""
    rng = random.Random(7)
    alphabet = string.ascii_letters + " -"
    for _ in range(300):
        text = "".join(rng.choice(alphabet) for _ in range(rng.randint(1, 14)))
        got = chunks(text, MERGES)
        assert "".join(s for _, s in got) == text
        letter = rng.choice(string.ascii_lowercase)
        assert count_from_chunks(text, letter, MERGES) == text.lower().count(letter)


def test_ids_alone_miss_letters_inside_longer_chunks():
    """The lesson's point: reading ids as letters undercounts whenever a letter sits in a bigger chunk."""
    assert count_ids_that_are_the_letter("strawberry", "r", MERGES) < count_from_chunks("strawberry", "r", MERGES)
    assert count_ids_that_are_the_letter("banana", "a", MERGES) == 1 and count_from_chunks("banana", "a", MERGES) == 3


@pytest.mark.parametrize("separator", [SEPARATOR, " "])
def test_spelled_out_words_have_one_letter_per_chunk(separator):
    """Guards the claim 'spell it out and every letter becomes a chunk of its own': in the toy, no chunk of a
    hyphen- or space-separated word holds two letters."""
    rng = random.Random(11)
    for _ in range(200):
        word = "".join(rng.choice(string.ascii_lowercase) for _ in range(rng.randint(1, MAX_WORD)))
        spelled = separator.join(word)
        assert max(sum(c.isalpha() for c in s) for _, s in chunks(spelled, MERGES)) == 1
    assert spelled_out("abc") == "a-b-c"


def test_non_ascii_text_is_refused():
    """Guards a silent wrong answer: a multi-byte character splits into bytes whose spellings cannot be shown."""
    with pytest.raises(ValueError):
        chunks("café", MERGES)


def test_more_merges_never_make_more_chunks():
    """The widget's 'merges used' slider: each extra merge can only join pieces, so the chunk count never rises
    and at zero merges every character is a chunk."""
    word = "unbelievable"
    sizes = [len(chunks(word, MERGES[:k])) for k in (0, 1, 10, 100, 300, 512)]
    assert sizes == sorted(sizes, reverse=True) and sizes[0] == len(word)


def test_capitals_and_a_leading_space_change_the_ids():
    """Guards the lesson's note that the same letters can reach the model as different ids."""
    plain = [i for i, _ in chunks("strawberry", MERGES)]
    assert [i for i, _ in chunks(" strawberry", MERGES)] != plain
    assert [i for i, _ in chunks("STRAWBERRY", MERGES)] != plain
    assert len(chunks("STRAWBERRY", MERGES)) == 10


def test_count_inside_ignores_case():
    assert count_inside("Err", "r") == 2 and count_inside("er", "R") == 1


def test_pieces_are_what_the_tokenizer_cuts_first():
    """The widget runs PIECE as a JavaScript regex; a hyphen is its own piece, so a spelled-out word never merges."""
    assert pieces("s-t") == ["s", "-", "t"]


# Cuts and ids the lesson prints for the published tokenizers (tiktoken 0.14.0, run 2026-10-19). "|" marks a cut.
CUTS = {
    "strawberry": ("str|aw|berry", "st|raw|berry"),
    "Strawberry": ("Str|aw|berry", "Str|aw|berry"),
    "STRAWBERRY": ("ST|RAW|B|ERRY", "ST|RAW|B|ERRY"),
    " strawberry": (" strawberry", " strawberry"),
    "banana": ("banana", "banana"),
    "Mississippi": ("Miss|issippi", "Miss|issippi"),
    "tokenizer": ("tokenizer", "token|izer"),
    "unbelievable": ("un|belie|vable", "un|bel|ievable"),
}
IDS = {
    "strawberry": ([496, 675, 15717], [302, 1618, 19772]),
    "Strawberry": ([2645, 675, 15717], [3504, 1134, 19772]),
    "STRAWBERRY": ([790, 22275, 33, 73995], [1117, 46176, 33, 132354]),
    " strawberry": ([73700], [101830]),
}


def test_published_tokenizer_numbers_in_the_lesson():
    """Guards the lesson's table, figures and quiz against a library or vocabulary that moves: every cut and id the
    lesson prints for cl100k_base and o200k_base, and the id shared by the three r chunks of the spelled-out word.
    Skipped when tiktoken or its vocabulary files are not available."""
    tiktoken = pytest.importorskip("tiktoken")
    try:
        encs = [tiktoken.get_encoding(name) for name in ("cl100k_base", "o200k_base")]
    except Exception as err:  # the first use downloads the vocabulary
        pytest.skip(f"vocabulary not available: {type(err).__name__}")
    assert [(e.name, e.n_vocab) for e in encs] == [("cl100k_base", 100277), ("o200k_base", 200019)]
    for text, want in CUTS.items():
        for enc, cut in zip(encs, want):
            assert "|".join(enc.decode([i]) for i in enc.encode(text)) == cut, (text, enc.name)
    for text, want in IDS.items():
        for enc, ids in zip(encs, want):
            assert enc.encode(text) == ids, (text, enc.name)
    for enc, r_id in zip(encs, (3880, 6335)):
        spelled = enc.encode(SEPARATOR.join("strawberry"))
        assert len(spelled) == 10 and [i for i in spelled if i == r_id] == [r_id] * 3
        assert [enc.decode([i]) for i in spelled][2] == "-r"