import pytest

from prompts import PROMPTS, true_sentence
from toy_lm import (IDK, answer, lines_of, load_corpus, model_with_extra, next_word_probs, ranked, smallest_k,
                    tokenize, train, with_temperature)

LINES = lines_of(load_corpus())
PERU = PROMPTS[0]
SENTENCE = true_sentence(PERU)


def probs_at(k, prompt=PERU["prompt"]):
    return next_word_probs(model_with_extra(LINES, SENTENCE, k), prompt)


@pytest.mark.parametrize("prompt", ["The capital of Peru is", "The capital of France is", "", "Peru", "zzz qqq", "is", "The capital of"])
@pytest.mark.parametrize("k", [0, 1, 7])
def test_probabilities_add_up_to_one(prompt, k):
    """Guards a backoff that leaks or double-counts probability (an unseen context, a one-word prompt, an empty one)."""
    probs = probs_at(k, prompt)
    assert sum(probs.values()) == pytest.approx(1, abs=1e-12)
    assert all(p >= 0 for p in probs.values())


@pytest.mark.parametrize("temperature", [0.1, 0.5, 1.0, 2.0])
def test_temperature_keeps_the_total_at_one(temperature):
    assert sum(with_temperature(probs_at(0), temperature).values()) == pytest.approx(1, abs=1e-12)


def test_thin_corpus_gives_a_fluent_wrong_answer():
    """The corpus never says what the capital of Peru, Laos or Malta is, yet the toy completes the sentence with
    a capital it has read about, and gives the true word no probability at all."""
    model = train(LINES)
    for item in PROMPTS:
        probs = next_word_probs(model, item["prompt"])
        top, p = ranked(probs)[0]
        assert top == "Paris" and top != item["truth"]
        assert item["truth"] not in model["vocab"] and probs.get(item["truth"], 0) == 0
        assert p == pytest.approx(0.1049, abs=5e-5)
        assert answer(probs) == "Paris"


def test_the_held_out_facts_are_not_in_the_corpus_in_any_wording():
    """Guards a corpus edit that quietly teaches the answer, which would make the demonstration dishonest."""
    text = load_corpus()
    for item in PROMPTS:
        assert item["truth"] not in text
        assert f"{item['name']} is" not in text            # no pair like ('Peru', 'is') to match the prompt
    assert ("Peru", "is") not in train(LINES)["pairs"]


def test_top_probability_matches_a_hand_calculation():
    """'is' is followed by 31 words, 19 of them different (30 capitals in 18 kinds, plus 'a' once). The trust in that
    context is 31/(31+19) = 0.62. Paris follows 'is' 5 times and the vocabulary has 77 words."""
    expected = 0.62 * 5 / 31 + 0.38 / 77
    assert next_word_probs(train(LINES), PERU["prompt"])["Paris"] == pytest.approx(expected, abs=1e-12)


def test_one_true_sentence_hand_calculation():
    """After one copy the pair ('Peru', 'is') was seen once, so its trust is 1/(1+1) = 0.5. The word 'is' now has 32
    followers, 20 of them different, so its trust is 32/52; Lima follows it once, and the vocabulary has 78 words."""
    single = (32 / 52) * (1 / 32) + (20 / 52) * (1 / 78)
    assert probs_at(1)["Lima"] == pytest.approx(0.5 + 0.5 * single, abs=1e-12)


def test_true_examples_take_over_and_the_code_finds_the_smallest_k():
    assert [ranked(probs_at(k))[0][0] for k in range(0, 4)] == ["Paris", "Lima", "Lima", "Lima"]
    assert smallest_k(LINES, PERU["prompt"], SENTENCE, PERU["truth"]) == 1
    for item in PROMPTS:
        assert smallest_k(LINES, item["prompt"], true_sentence(item), item["truth"]) == 1
    # Lima's share climbs with every copy
    shares = [probs_at(k)["Lima"] for k in range(1, 12)]
    assert shares == sorted(shares) and shares[0] == pytest.approx(0.5121, abs=5e-5)


@pytest.mark.parametrize("threshold, k", [(0.5, 1), (0.8, 4), (0.9, 8)])
def test_smallest_k_for_a_confidence_target(threshold, k):
    """Guards an off-by-one in the search: k-1 copies must fall short and k copies must reach the target."""
    assert smallest_k(LINES, PERU["prompt"], SENTENCE, PERU["truth"], threshold) == k
    assert ranked(probs_at(k))[0][1] >= threshold > ranked(probs_at(k - 1))[0][1]


def test_smallest_k_gives_up_at_the_limit():
    """Guards an off-by-one at the limit: 4 copies are needed at the 80% bar, so limit=4 finds them and limit=3 does not."""
    assert smallest_k(LINES, PERU["prompt"], SENTENCE, PERU["truth"], 0.8, limit=4) == 4
    assert smallest_k(LINES, PERU["prompt"], SENTENCE, PERU["truth"], 0.8, limit=3) is None


def test_abstaining_when_the_top_word_is_below_the_threshold():
    probs = probs_at(0)
    assert answer(probs, 0.5) == IDK
    assert answer(probs, 0.11) == IDK
    assert answer(probs, 0.10) == "Paris"                  # 10.5% clears a 10% bar: a low bar does not help
    assert answer(probs_at(1), 0.5) == "Lima"              # the same bar lets the right answer through
    assert answer(next_word_probs(train(LINES), "The capital of France is"), 0.5) == "Paris"   # and a well-read fact


def test_the_toy_gets_a_fact_it_has_read_often_right():
    """The control: the corpus states France's capital five times, so the same code answers it with 85%."""
    probs = next_word_probs(train(LINES), "The capital of France is")
    assert ranked(probs)[0][0] == "Paris" and ranked(probs)[0][1] == pytest.approx(0.8508, abs=5e-5)


def test_lower_temperature_makes_the_wrong_answer_surer_not_true():
    base = probs_at(0)
    shares = []
    for t in (2.0, 1.0, 0.5, 0.3, 0.1):
        scaled = with_temperature(base, t)
        top, p = ranked(scaled)[0]
        assert top == "Paris" and scaled.get("Lima", 0) == 0
        shares.append(p)
    assert shares == sorted(shares) and shares[-1] > 0.98
    # a 50% bar that makes the toy say "I don't know" at temperature 1 lets Paris through at 0.3
    assert answer(with_temperature(base, 1.0), 0.5) == IDK
    assert answer(with_temperature(base, 0.3), 0.5) == "Paris"


def test_temperature_one_changes_nothing():
    base = probs_at(0)
    scaled = with_temperature(base, 1.0)
    assert all(scaled[w] == pytest.approx(base[w], abs=1e-15) for w in base)


def test_a_tie_goes_to_the_word_that_sorts_first_whatever_the_corpus_order():
    """Guards the vocabulary order that ties rely on: it is sorted, so equal words rank alphabetically in either corpus order."""
    for lines in (["x b", "x a"], ["x a", "x b"]):
        probs = next_word_probs(train(lines), "x")
        assert probs["a"] == probs["b"]
        assert [w for w, _ in ranked(probs)] == ["a", "b", "x"]


def test_an_empty_corpus_has_nothing_to_say():
    assert next_word_probs(train([]), "anything") == {}
    assert answer({}, 0.0) == IDK


def test_tokenizer_keeps_ascii_words_and_marks_only():
    assert tokenize("Peru's capital, Lima! Café 2026?") == ["Peru's", "capital", ",", "Lima", "!", "Caf", "2026", "?"]
    assert lines_of("a\n\n  \nb\n") == ["a", "b"]