import pytest

from prompts import PROMPTS
from toy_gen import (END, ORDER, START, chances, generate, lines_of, load_corpus, pick, positions_read, reply_of,
                     seed_state, seed_step, tokenize, train)

LINES = lines_of(load_corpus())
MODEL = train(LINES)
TEA = PROMPTS[0]["prompt"]
GREEDY = {
    "Tea": "Boil some water. Put a tea bag in a cup. Pour the water over the bag. Wait three minutes. Pull the bag out. Enjoy.",
    "Egg": "Put the egg in a pan. Cover it with cold water. Bring the water to a boil. Wait ten minutes. Scoop the egg out. Enjoy.",
    "Toast": "Slice some bread. Put the slice in a toaster. Press the lever down. Wait two minutes. Lift out the toast. Add butter. Enjoy.",
    "Plant": "Check the soil with a finger. If the soil is dry, add water. Pour slowly until water drains out. Empty the dish under the pot. Check again next week.",
}


def test_every_question_opens_at_least_two_corpus_lines():
    """Guards a widget question the corpus cannot answer: the first pass of each would fall back to guessing from all tokens."""
    for item in PROMPTS:
        assert sum(line.startswith(item["prompt"]) for line in LINES) >= 2, item["prompt"]


@pytest.mark.parametrize("item", PROMPTS, ids=lambda item: item["name"])
def test_greedy_reply_is_the_most_common_line_for_that_question(item):
    """Guards a pick that is not the likeliest token: the repeatable demo must write the line the corpus repeats most."""
    passes = generate(MODEL, item["prompt"], 40)
    assert reply_of(passes) == GREEDY[item["name"]]
    assert passes[-1]["token"] == END
    assert item["prompt"] + " " + GREEDY[item["name"]] in LINES


@pytest.mark.parametrize("item", PROMPTS, ids=lambda item: item["name"])
def test_one_token_per_pass(item):
    """Guards a loop that adds two tokens, or none, in a pass: each pass adds exactly one token to the text, so
    the passes equal the reply tokens plus the end mark, and every pass starts from the text the last one left."""
    passes = generate(MODEL, item["prompt"], 40)
    written = [p["token"] for p in passes if p["token"] != END]
    assert len(passes) == len(written) + 1 == len(tokenize(reply_of(passes))) + 1
    text = tokenize(item["prompt"])
    for p in passes:
        assert p["before"] == text
        if p["token"] != END:
            text = text + [p["token"]]


@pytest.mark.parametrize("item", PROMPTS, ids=lambda item: item["name"])
def test_chances_add_up_to_one_and_are_ranked(item):
    for p in generate(MODEL, item["prompt"], 40):
        assert sum(count for _, count in p["ranked"]) == p["total"]
        assert sum(count / p["total"] for _, count in p["ranked"]) == pytest.approx(1, abs=1e-12)
        keys = [(-count, token) for token, count in p["ranked"]]
        assert keys == sorted(keys)


def test_first_pass_on_tea_splits_four_to_one():
    """The numbers the lesson opens with: of the 5 lines that start with the tea question, 4 go on with Boil and 1 with Heat."""
    looked_at, ranked, total = chances(MODEL, tokenize(TEA))
    assert looked_at == ("make", "tea", "?")
    assert ranked == [("Boil", 4), ("Heat", 1)] and total == 5


def test_first_three_passes_match_the_hand_worked_table():
    """Guards the lesson's hand-worked table drifting from the code: the three rows, the 5 lines behind the first, and the
    text growing from 6 to 9 tokens."""
    passes = generate(MODEL, TEA, 3)
    assert [(p["looked_at"], p["ranked"], p["total"], p["token"]) for p in passes] == [
        (["make", "tea", "?"], [("Boil", 4), ("Heat", 1)], 5, "Boil"),
        (["tea", "?", "Boil"], [("some", 4)], 4, "some"),
        (["?", "Boil", "some"], [("water", 4)], 4, "water"),
    ]
    assert [len(p["before"]) for p in passes] == [6, 7, 8]
    assert sum(line.startswith(TEA) for line in LINES) == 5
    assert positions_read(6, 3, False) == 6 + 7 + 8 and positions_read(6, 3, True) == 6 + 1 + 1


def test_a_pass_reads_only_the_last_three_tokens():
    """Guards a model that quietly reads more than three tokens: text before the last three changes nothing."""
    tokens = tokenize(TEA) + ["Boil", "some", "water", ".", "Put"]
    base = chances(MODEL, tokens)
    assert base[0] == ("water", ".", "Put") and len(base[0]) == ORDER
    assert chances(MODEL, ["Elephant", "?"] + tokens) == base
    assert chances(MODEL, tokens[3:]) == base


def test_unseen_text_falls_back_to_shorter_runs():
    """Guards a crash or an empty list when the last three tokens were never seen: look at two, then one, then all."""
    assert chances(MODEL, ["zebra", "make", "tea", "?"])[0] == ("make", "tea", "?")
    assert chances(MODEL, ["zebra", "zebra", "tea", "?"])[0] == ("tea", "?")        # "zebra tea ?" never occurs, "tea ?" does
    assert chances(MODEL, ["zebra", "zebra", "zebra", "water"])[0] == ("water",)
    looked_at, ranked, total = chances(MODEL, ["zebra"])
    assert looked_at == () and total == sum(count for _, count in ranked) and ranked[0][0] == "."
    assert chances(train([]), ["a"]) == ((), [], 0)


def test_start_marks_pad_the_beginning_of_a_line():
    """Guards the first token of a line being counted with no context: two tokens of text still find the lines that begin with them."""
    model = train(["Hello there friend."])
    assert chances(model, [])[1] == [("Hello", 1)] and chances(model, ["Hello"])[1] == [("there", 1)]
    assert START not in [token for k in range(ORDER + 1) for followers in model[k].values() for token in followers]


@pytest.mark.parametrize("item", PROMPTS, ids=lambda item: item["name"])
def test_greedy_runs_are_repeatable_and_stop_at_the_end_mark(item):
    assert generate(MODEL, item["prompt"], 40) == generate(MODEL, item["prompt"], 40)
    assert len(generate(MODEL, item["prompt"], 40)) < 40
    assert len(generate(MODEL, item["prompt"], 1000)) == len(generate(MODEL, item["prompt"], 40))


@pytest.mark.parametrize("cap", [1, 4, 10, 30])
def test_length_cap_stops_the_reply_without_an_end_mark(cap):
    """Guards a cap that is off by one: a cap of N allows N passes, and the reply then stops mid-sentence."""
    passes = generate(MODEL, TEA, cap)
    assert len(passes) == cap and all(p["token"] != END for p in passes)


def test_seed_generator_matches_the_published_check_value():
    """Park-Miller with factor 48271 (minstd_rand): from state 1, the 10000th state is 399268537. Guards a changed constant,
    which would still give plausible numbers but different picks in JavaScript and Python."""
    state = 1
    for _ in range(10000):
        state, u = seed_step(state)
    assert state == 399268537
    assert seed_state(0) == 1 and seed_state(1) == 2 and seed_state(2147483645) == 2147483646 and seed_state(2147483646) == 1
    assert seed_step(seed_state(1)) == (96542, 96541 / 2147483646)


def test_pick_by_chance_gives_each_token_its_share():
    """Guards an off-by-one in the walk: with one draw per unit of the total, each token is picked exactly count times."""
    ranked, total = [("a", 3), ("b", 2), ("c", 1)], 6
    drawn = [pick(ranked, total, (i + 0.5) / total) for i in range(total)]
    assert drawn == ["a", "a", "a", "b", "b", "c"]
    assert pick(ranked, total) == "a" and pick(ranked, total, 0.0) == "a" and pick(ranked, total, 0.9999999) == "c"


def test_seeded_runs_repeat_and_differ_between_seeds():
    assert generate(MODEL, TEA, 40, seed=7) == generate(MODEL, TEA, 40, seed=7)
    replies = {reply_of(generate(MODEL, TEA, 40, seed=s)) for s in range(1, 21)}
    assert GREEDY["Tea"] in replies and len(replies) >= 3
    assert reply_of(generate(MODEL, TEA, 40, seed=7)) != reply_of(generate(MODEL, TEA, 40, seed=1))


def test_seeded_runs_stay_on_the_question_and_end():
    """Guards a corpus where two questions share three-token runs: a draw could then wander into another question's recipe
    (a tea reply that takes out the toast). Every token must come from a line that starts with the same question."""
    for item in PROMPTS:
        own = {t for line in LINES if line.startswith(item["prompt"]) for t in tokenize(line)} | {END}
        for seed in range(1, 301):
            passes = generate(MODEL, item["prompt"], 40, seed)
            assert passes[-1]["token"] == END and {p["token"] for p in passes} <= own, (item["name"], seed)


def simulate_positions(prompt_tokens, new_tokens, cached):
    """Run the loop on paper: each pass runs the network over the text so far, or only over what the cache lacks."""
    done = total = 0
    for i in range(new_tokens):
        length = prompt_tokens + i
        total += length - done if cached else length
        done = length
    return total


def test_positions_read_matches_counting_pass_by_pass():
    """Guards the closed form behind the cache numbers: run the loop on paper and compare."""
    for prompt_tokens in (1, 6, 7):
        for new in range(0, 12):
            for cached in (False, True):
                assert positions_read(prompt_tokens, new, cached) == simulate_positions(prompt_tokens, new, cached)
    assert positions_read(6, 1, False) == positions_read(6, 1, True) == 6
    assert (positions_read(6, 31, False), positions_read(6, 31, True)) == (651, 36)


def test_reply_text_attaches_marks_to_the_word_before():
    assert reply_of([{"token": t} for t in ["Hi", ",", "you", ".", END, "x"]]) == "Hi, you."
    assert reply_of([]) == ""