"""Tests for sampler.py. Each one names the mistake it would catch.

Run: PYTHONDONTWRITEBYTECODE=1 python3 -m pytest -q -p no:cacheprovider .
"""
from collections import Counter
from math import isclose

import pytest

import sampler as s

TOY = s.toy
LINES = TOY.lines_of(TOY.load_corpus())
MODEL = TOY.train(LINES)
TEA = "How do I make tea?"
PLANT = "How do I water a plant?"
WAIT = s.as_chances(*s.row_of(MODEL, ["Wait"]))        # the first row the widget opens on


def share(chances):
    return {token: chance for token, chance in chances}


def test_rows_are_lesson_4_numbers():
    """Lesson 5 must show the very chances lesson 4 showed: pass 1 of the tea question and pass 11."""
    assert s.row_of(MODEL, ["make", "tea", "?"]) == ([("Boil", 4), ("Heat", 1)], 5)
    assert s.row_of(MODEL, ["bag", "in", "a"]) == ([("cup", 3), ("mug", 2)], 5)
    greedy = TOY.generate(MODEL, TEA, 40)
    assert greedy[0]["ranked"] == [("Boil", 4), ("Heat", 1)]
    assert greedy[10]["ranked"] == [("cup", 3), ("mug", 2)]


def test_every_widget_row_exists_and_has_followers():
    """A row naming a context the toy never saw would crash the widget at start."""
    assert len({row["id"] for row in s.ROWS}) == len(s.ROWS)
    for row in s.ROWS:
        ranked, total = s.row_of(MODEL, row["context"])
        assert total == sum(count for _, count in ranked) > 0


def test_wait_row_is_what_the_lesson_says():
    """The lesson's bag of 13 marbles: three 4, ten 3, two 3, and one each for four, one, six."""
    ranked, total = s.row_of(MODEL, ["Wait"])
    assert total == 13
    assert ranked == [("three", 4), ("ten", 3), ("two", 3), ("four", 1), ("one", 1), ("six", 1)]


def test_generator_is_lesson_4s_park_miller():
    """Changing the factor or the modulus would change every draw in lesson 4 and here."""
    state = 1
    for _ in range(10000):
        state, _ = TOY.seed_step(state)
    assert state == 399268537


def test_choose_agrees_with_lesson_4_pick():
    """At temperature 1 and p = 1 the draw must be lesson 4's pick, or the two lessons would disagree on one seed."""
    for row in s.ROWS:
        ranked, total = s.row_of(MODEL, row["context"])
        chances = s.as_chances(ranked, total)
        state = s.start_state(7)
        for _ in range(2000):
            state, u = TOY.seed_step(state)
            assert s.choose(chances, u) == TOY.pick(ranked, total, u)


def test_choose_walks_the_slices_in_order():
    """u picks by position: the first slice starts at 0, and u past the rounded-off total takes the last token."""
    chances = [("a", 0.5), ("b", 0.25), ("c", 0.25)]
    assert [s.choose(chances, u) for u in (0.0, 0.49, 0.5, 0.74, 0.75, 0.99)] == ["a", "a", "b", "b", "c", "c"]
    assert s.choose([("a", 0.3), ("b", 0.3)], 0.7) == "b"


def test_temperature_one_changes_nothing():
    assert all(isclose(a, b) for (_, a), (_, b) in zip(WAIT, s.with_temperature(WAIT, 1.0)))


def test_temperature_below_one_sharpens_and_above_one_flattens():
    """The direction is the whole point: low favours the top token, high evens the tokens out."""
    base = WAIT[0][1]
    assert s.with_temperature(WAIT, 0.5)[0][1] > base > s.with_temperature(WAIT, 2.0)[0][1]
    assert s.with_temperature(WAIT, 0.05)[0][1] > 0.99
    spread = lambda chances: max(c for _, c in chances) - min(c for _, c in chances)
    assert spread(s.with_temperature(WAIT, 50.0)) < 0.01 < spread(WAIT)


def test_temperature_keeps_total_one_and_order():
    for t in (0.05, 0.2, 0.5, 1.0, 2.0):
        out = s.with_temperature(WAIT, t)
        assert isclose(sum(c for _, c in out), 1.0)
        assert [tok for tok, _ in out] == [tok for tok, _ in WAIT]


def test_temperature_is_lesson_12s_power_rule_and_softmax_of_scores_over_t():
    """Dividing the scores by T before softmax equals raising the chances to 1/T and rescaling (the Go deeper claim)."""
    from math import exp, log
    for t in (0.3, 0.7, 1.5):
        scores = [log(c) for _, c in WAIT]                       # scores whose softmax is the chances
        exps = [exp(x / t) for x in scores]
        soft = [e / sum(exps) for e in exps]
        assert all(isclose(a, c, rel_tol=1e-12) for a, (_, c) in zip(soft, s.with_temperature(WAIT, t)))


@pytest.mark.parametrize("bad", [0, 0.0, -1])
def test_temperature_zero_or_below_is_refused(bad):
    """Zero would divide by zero; the toy asks for the greedy choice some other way."""
    with pytest.raises(ValueError):
        s.with_temperature(WAIT, bad)


def test_top_p_keeps_the_smallest_top_set_and_rescales():
    kept = s.top_p(WAIT, 0.7)
    assert [t for t, _ in kept] == ["three", "ten", "two"]                    # 10 of 13 is 76.9%, the first total at or past 0.7
    assert [round(c, 3) for _, c in kept] == [0.4, 0.3, 0.3]
    assert isclose(sum(c for _, c in kept), 1.0)
    assert [t for t, _ in s.top_p(WAIT, 0.5)] == ["three", "ten"]


def test_top_p_one_keeps_everything_and_tiny_p_keeps_one():
    assert [t for t, _ in s.top_p(WAIT, 1.0)] == [t for t, _ in WAIT]
    assert s.top_p(WAIT, 0.001) == [("three", 1.0)]
    cold = s.with_temperature(WAIT, 0.05)                                  # four, one, six hold under 1e-9 between them
    assert len(s.top_p(cold, 1.0)) == 6 and len(s.top_p(cold, 0.9999)) < 6


def test_top_p_tie_goes_to_the_token_that_sorts_first():
    """At p = 0.8 only one of four, one and six fits, and which one is a rule, not luck: the one that sorts first."""
    assert [t for t, _ in s.top_p(WAIT, 0.8)] == ["three", "ten", "two", "four"]


def test_top_p_rounding_does_not_pull_in_one_more_token():
    """0.7 + 0.2 is 0.8999999999999999 in floating point; a cut at exactly 0.9 must still stop there."""
    assert [t for t, _ in s.top_p([("a", 0.7), ("b", 0.2), ("c", 0.1)], 0.9)] == ["a", "b"]


def test_reshape_applies_temperature_then_top_p():
    """The order matters: cutting first and sharpening after keeps three tokens at 0.2 and 0.7; this order keeps two."""
    assert [t for t, _ in s.reshape(WAIT, 0.2, 0.7)] == ["three", "ten"]
    assert [t for t, _ in s.with_temperature(s.top_p(WAIT, 0.7), 0.2)] == ["three", "ten", "two"]
    assert [t for t, _ in s.reshape(WAIT, 0.1, 0.7)] == ["three"]
    assert [t for t, _ in s.reshape(WAIT, 2.0, 0.7)] == ["three", "ten", "two", "four"]


def test_same_seed_same_draws_different_seed_different_draws():
    """The seed is what makes the toy repeatable; without it the lesson's 'same seed, same draws' would be false."""
    a, _ = s.draw(WAIT, s.start_state(1))
    b, _ = s.draw(WAIT, s.start_state(1))
    c, _ = s.draw(WAIT, s.start_state(2))
    assert a == b and a != c


def test_draw_goes_on_from_the_state_it_returns():
    """Pressing Draw 20 times again must continue the stream, not repeat the same 20 picks."""
    state = s.start_state(1)
    first, state = s.draw(WAIT, state)
    second, state = s.draw(WAIT, state)
    both, _ = s.draw(WAIT, s.start_state(1), 40)
    assert first + second == both and first != second


def test_the_opening_draws_the_lesson_quotes():
    """Seed 1, temperature 1, no cut: the 20 picks and counts the widget opens on."""
    picks, _ = s.draw(WAIT, s.start_state(1))
    tokens = [t for t, _ in WAIT]
    assert s.tally(picks, tokens) == [7, 4, 3, 2, 1, 3]
    assert len(picks) == 20


def test_small_seeds_do_not_always_start_with_the_top_token():
    """Without the warm-up a seed of 1 makes a first number near 0.00004 and the first draw is always the top token."""
    raw = TOY.seed_step(TOY.seed_state(1))[1]
    assert raw < 1e-4
    heat = lambda start: sum(s.generate(MODEL, TEA, start(seed), cap=1)[0] == ["Heat"] for seed in range(1000))
    assert heat(TOY.seed_state) == 0                                      # the lesson's 0 of 1,000 without the warm-up
    assert 150 < heat(s.start_state) < 250                                # Heat has a 20% chance on the first pass


def test_draws_follow_the_chances():
    """Frequency check: with a fixed seed, 100,000 draws land within a point of the chances (never flaky)."""
    for temperature, p in ((1.0, 1.0), (0.5, 1.0), (2.0, 0.8)):
        chances = s.reshape(WAIT, temperature, p)
        picks, _ = s.draw(chances, s.start_state(1), 100000)
        seen = Counter(picks)
        for token, chance in chances:
            assert abs(seen[token] / 100000 - chance) < 0.01
        assert set(seen) == {token for token, _ in chances}               # a token top-p cut never comes up


def test_low_temperature_replies_follow_the_greedy_reply():
    """At a very low temperature the toy writes lesson 4's greedy reply, for any seed, where the top token is unique."""
    greedy = [p["token"] for p in TOY.generate(MODEL, TEA, 40)]
    for seed in (1, 2, 3):
        reply, _ = s.generate(MODEL, TEA, s.start_state(seed), temperature=0.01)
        assert reply == greedy


def test_generate_is_repeatable_and_regenerate_continues_the_stream():
    one, state = s.generate(MODEL, PLANT, s.start_state(1))
    again, _ = s.generate(MODEL, PLANT, s.start_state(1))
    two, _ = s.generate(MODEL, PLANT, state)
    assert one == again
    assert s.replies(MODEL, PLANT, 2, 1) == [one, two]


def test_reply_counts_the_lesson_quotes():
    """1,000 tea replies from one seed: how many different ones, and how many are not among the toy's own lines."""
    known = s.corpus_replies(LINES, TEA)
    assert len(known) == 4
    texts = Counter(s.reply_text(r) for r in s.replies(MODEL, TEA, 1000, 1))
    assert (len(texts), sum(1 for t in texts if t not in known)) == (16, 12)
    cool = Counter(s.reply_text(r) for r in s.replies(MODEL, TEA, 1000, 1, temperature=0.05))
    assert len(cool) == 2 and cool.most_common(1)[0][1] == 999 and set(cool) <= known
    assert len({s.reply_text(r) for r in s.replies(MODEL, TEA, 1000, 1, temperature=0.01)}) == 1
    cut = {s.reply_text(r) for r in s.replies(MODEL, TEA, 1000, 1, p=0.7)}            # top-p 0.7 in 'How many different replies'
    assert (len(cut), len(cut - known)) == (4, 2)


def test_plant_runs_part_ways_at_dry_or_wet():
    """The branch the lesson shows: seed 1's first and sixth replies share 11 tokens, then one says dry and one wet."""
    runs = s.replies(MODEL, PLANT, 6, 1)
    dry, wet = runs[0], runs[5]
    fork = s.first_split(dry, wet)
    assert fork == 11
    assert (dry[fork], wet[fork]) == ("dry", "wet")
    assert (len(dry) - fork - 1, len(wet) - fork - 1) == (24, 12)           # tokens from dry or wet on, end mark not counted
    assert s.row_of(MODEL, ["the", "soil", "is"]) == ([("dry", 3), ("wet", 1)], 4)


def test_first_split():
    assert s.first_split(["a", "b"], ["a", "c"]) == 1
    assert s.first_split(["a"], ["a", "b"]) == 1
    assert s.first_split(["a", "b"], ["a", "b"]) is None