"""Guards for safety_filter.py: the data, the training, the counting at every threshold, the rare-harm arithmetic, and the failures the lesson shows."""
import math
import re
import textwrap
from fractions import Fraction
from pathlib import Path

import pytest

import make_figures
import safety_filter as sf

LESSON = Path(__file__).resolve().parent.parent
SAME = 68                                           # the first 68 test messages use wordings the training messages had; the other 52 do not


@pytest.fixture(scope="module")
def r():
    return sf.build()


def test_data_has_the_sizes_the_lesson_names_and_the_test_never_repeats_a_training_message(r):
    """Guards a test set that quietly overlaps the training set, which would make the filter look better than it is."""
    labels = r["labels"]
    assert [len(r["train"]), sum(y for _, y in r["train"])] == [300, 100]
    assert [len(r["test"]), sum(labels), sum(labels[SAME:]), len(labels) - SAME] == [120, 40, 16, 52]
    assert not {t for t, _ in r["train"]} & {t for t, _ in r["test"]}
    assert sf.make_data()[1] == r["test"]                                # fixed seed: the same messages every run
    assert not set(sf.UNSEEN_HARMFUL) & set(sf.HARMFUL) and not set(sf.UNSEEN_HARMLESS) & set(sf.HARMLESS)       # "never seen" means no template in common


def test_training_loss_starts_at_ln_2_and_never_rises(r):
    """Guards a training loop that stopped learning: a filter with all weights zero is a coin flip, whose loss is ln 2."""
    losses = r["losses"]
    assert losses[0] == pytest.approx(math.log(2)) and losses[-1] < 0.05
    assert all(b <= a + 1e-12 for a, b in zip(losses, losses[1:]))


def test_counts_add_up_and_move_one_way_as_the_threshold_rises(r):
    """Guards the sweep behind the lesson's curves: at every threshold the four counts cover all 120 messages, a higher threshold
    only ever misses more and falsely blocks fewer, and the rates at 0.50 are the ones the lesson names."""
    table = r["sweep"]
    for p in range(1, 100):
        caught, missed, false_alarm, passed = table[p]
        assert (caught + missed, false_alarm + passed) == (40, 80)
    assert all(table[p + 1][1] >= table[p][1] and table[p + 1][2] <= table[p][2] for p in range(1, 99))
    miss, alarm, precision = sf.rates(table[50])
    assert (table[50], miss, alarm) == ((35, 5, 1, 79), 5 / 40, 1 / 80) and round(100 * precision) == 97


def test_no_threshold_is_right_for_every_message(r):
    """Guards the lesson's main claim: no threshold has both mistakes at zero, the two rates cross, and each end of the range
    trades one mistake for the other."""
    table = r["sweep"]
    assert not [p for p, c in table.items() if c[1] == 0 and c[2] == 0]
    assert 1 < sf.crossing(table) < 99
    assert table[1][1] == 0 < table[1][2] and table[99][2] == 0 < table[99][1]
    assert min(c[1] + c[2] for c in table.values()) == 6


def test_a_message_scoring_exactly_the_threshold_is_blocked():
    assert sf.counts([1, 0], [0.4, 0.4], 40) == (1, 0, 1, 0) and sf.counts([1, 0], [0.399, 0.399], 40) == (0, 1, 0, 1)


def test_rare_harm_arithmetic_matches_exact_fractions_and_flips_the_precision(r):
    """Guards the base-rate numbers: at 1 harmful in 100 the same rates block mostly harmless messages. The widget and the lesson
    use whole numbers, rounding halves up, so check them against exact fractions at every threshold."""
    for p in range(1, 100):
        caught, missed, false_alarm, passed = r["sweep"][p]
        c, m, f, ok = sf.at_traffic(r["sweep"][p])
        assert (c + m, f + ok) == (100, 9900)
        assert abs(Fraction(c) - Fraction(100 * caught, 40)) <= Fraction(1, 2) and abs(Fraction(f) - Fraction(9900 * false_alarm, 80)) <= Fraction(1, 2)
    assert sf.at_traffic(r["sweep"][50]) == (88, 12, 124, 9776)
    on_test, on_traffic = sf.view(r["labels"], r["scores"], 50), sf.view(r["labels"], r["scores"], 50, True)
    assert on_test["precision"].endswith("(97%).") and on_traffic["precision"].endswith("(42%).")
    assert sf.share(88, 212) == 42 and sf.share(1, 2) == 50 and sf.share(1, 8) == 13        # a half rounds up


def test_the_three_examples_fail_or_pass_for_the_reasons_the_lesson_gives(r):
    """Guards the lesson's examples: a harmful message in new words is missed because its words carry no weight, a harmless
    message is blocked because it uses a word harmful messages use, and a harmless message with the word 'address' still passes."""
    test, scores, labels, vocab, w = r["test"], r["scores"], r["labels"], r["vocab"], r["w"]
    miss, alarm, passed = (sf.find(test, c) for c in (sf.MISSED_EXAMPLE, sf.FALSE_ALARM_EXAMPLE, sf.PASSED_EXAMPLE))
    assert (labels[miss], scores[miss] < 0.5, labels[alarm], scores[alarm] >= 0.5, labels[passed], scores[passed] < 0.5) == (1, True, 0, True, 0, True)
    known, unseen, _ = sf.explain(test[miss][0], r)
    assert {"w:drop", "w:ur", "w:digits"} <= set(unseen) and not {"w:drop", "w:ur", "w:digits"} & set(vocab)
    harmful_train = " ".join(t for t, y in r["train"] if y)
    assert "phone" in harmful_train and "battery" not in " ".join(t for t, _ in r["train"]) and float(w[vocab["w:phone"]]) > 0.4
    assert "address" in harmful_train and float(w[vocab["w:address"]]) > 0.4 and scores[passed] < 0.3
    mistakes = [i for i in range(len(test)) if (scores[i] >= 0.5) != bool(labels[i])]
    assert len(mistakes) == 6 and min(mistakes) >= SAME                   # every mistake is in wording the filter never saw


def test_explain_adds_up_to_the_score(r):
    """Guards the worked table in the lesson: bias plus the weights of a message's features, through the sigmoid, is its score."""
    for core in (sf.MISSED_EXAMPLE, sf.FALSE_ALARM_EXAMPLE, sf.PASSED_EXAMPLE):
        i = sf.find(r["test"], core)
        known, _, z = sf.explain(r["test"][i][0], r)
        assert round(float(sf.sigmoid(z)), 3) == r["scores"][i]
        assert z == pytest.approx(r["b"] + sum(v for _, v in known))


def test_the_toy_uses_only_numpy_and_the_standard_library_and_no_network():
    source = Path(sf.__file__).read_text()
    assert set(re.findall(r"^(?:from|import) (\w+)", source, flags=re.M)) == {"__future__", "re", "numpy"}


def test_every_python_block_in_the_lesson_is_a_verbatim_excerpt_of_the_source():
    """The lesson quotes this file. If the code changes and the lesson does not, the reader sees code that was never tested."""
    lesson = (LESSON / "lesson.md").read_text()
    source = Path(sf.__file__).read_text().splitlines()
    blocks = re.findall(r"^```python\n(.*?)^```", lesson, flags=re.S | re.M)
    assert blocks, "the lesson has no python blocks to check"
    for block in blocks:
        lines = block.rstrip("\n").splitlines()
        windows = (textwrap.dedent("\n".join(source[i:i + len(lines)])) for i in range(len(source) - len(lines) + 1))
        assert textwrap.dedent("\n".join(lines)) in windows, f"not in safety_filter.py:\n{block}"


def test_the_lesson_quotes_the_numbers_the_code_gives(r):
    """Guards prose that drifts from the code: the numbers in the opening paragraph, both threshold tables, the weight table, the
    worked examples, the by-hand arithmetic, the sweep text and the Check-yourself answers are recomputed here and must appear
    in lesson.md."""
    lesson = (LESSON / "lesson.md").read_text()
    sweep, train, test, scores, labels, w, vocab = (r[k] for k in ("sweep", "train", "test", "scores", "labels", "w", "vocab"))
    caught, missed, false_alarm, passed = sweep[50]
    harmful, harmless = caught + missed, false_alarm + passed
    caught_t, missed_t, false_t, _ = traffic = sf.at_traffic(sweep[50])
    blocked_t, harmless_t = caught_t + false_t, sf.TRAFFIC - sf.TRAFFIC_HARMFUL
    miss, alarm, _ = sf.rates(sweep[50])
    assert f"misses {missed} of {harmful} harmful test messages and blocks {false_alarm} of {harmless} harmless ones" in lesson
    assert f"would block about {traffic[0] + false_t} of every 10,000 messages, and about {false_t} of those would be harmless" in lesson
    assert f"row of {len(vocab)} numbers" in lesson and f"loss reaches {r['losses'][-1]:.4f}" in lesson
    for p in (1, 10, 25, 30, 40, 50, 60, 70, 90, 99):
        c = sweep[p]
        assert f"| 0.{p:02d} | {c[1]} | {c[2]} |" in lesson
    for p in (30, 40, 50, 60, 66, 70, 90):
        c, _, f, _ = sf.at_traffic(sweep[p])
        assert f"| 0.{p:02d} | {f:,} | {sf.share(c, c + f)}% |" in lesson
    words = lambda text: set(re.findall(r"[a-z0-9]+", text))
    in_train = lambda word, label: sum(1 for t, y in train if y == label and word in words(t))
    assert not (in_train("phone", 0) or in_train("address", 0) or in_train("the", 1)) and in_train("phone", 1) == in_train("address", 1)
    miss_i, alarm_i, pass_i = (sf.find(test, c) for c in (sf.MISSED_EXAMPLE, sf.FALSE_ALARM_EXAMPLE, sf.PASSED_EXAMPLE))
    wrong, unseen = [i for i in range(len(test)) if (scores[i] >= 0.5) != bool(labels[i])], len(test) - SAME
    fewest = [p for p, c in sweep.items() if c[1] + c[2] == min(c[1] + c[2] for c in sweep.values())]
    cross, half = sf.crossing(sweep), next(p for p in range(1, 100) if sf.at_traffic(sweep[p])[0] >= sf.at_traffic(sweep[p])[2])
    worked = []                                                     # the two worked tables, row by row
    for i in (miss_i, alarm_i):
        known, _, z = sf.explain(test[i][0], r)
        pairs = [v for n, v in known if n.startswith("c:")]
        parts = [r["b"], sum(pairs)] + [v for n, v in known if n.startswith("w:")]
        assert abs(sum(round(v, 3) for v in parts) - round(z, 3)) < 0.002       # the lesson says the shown parts add to the sum within 0.002
        worked += [f"| Bias | {r['b']:.3f} |", f"| Letter pairs ({len(pairs)}) | {sum(pairs):.3f} |", f"| Sum | {z:.3f} |", f"| Score | {scores[i]:.3f} |"]
        worked += [f'| Word "{n[2:]}" | {v:.3f} |' for n, v in known if n.startswith("w:")]
    known_p, _, z_p = sf.explain(test[pass_i][0], r)
    weight = lambda word: float(w[vocab["w:" + word]])
    quoted = [
        # the toy, the data and the by-hand arithmetic
        f"learned from {len(train)} made-up messages and is tested on {len(test)} others",
        f"There are {len(train)} training messages ({sum(y for _, y in train)} harmful and {sum(1 - y for _, y in train)} harmless) and {len(test)} test messages "
        f"({sum(labels)} harmful and {harmless} harmless)",
        f"{SAME} use wordings that training also had. The other {unseen} ({sum(labels[SAME:])} harmful and {unseen - sum(labels[SAME:])} harmless)",
        f"1 message in {len(test) // sum(labels)} is harmful", f"1 message in {sf.TRAFFIC // sf.TRAFFIC_HARMFUL} is harmful",
        f"Of {sf.TRAFFIC:,} messages, {sf.TRAFFIC_HARMFUL} are harmful and {harmless_t:,} are not",
        f"blocks about {caught_t} harmful messages and about {false_t} harmless ones. So {sf.share(false_t, blocked_t)}% of what it blocks is harmless, "
        f"although its false-alarm rate is a low {100 * alarm:.2f}%",
        f"Miss rate: {missed} missed of {harmful} harmful is {100 * miss:.1f}%", f"False-alarm rate: {false_alarm} blocked of {harmless} harmless is {100 * alarm:.2f}%",
        f"{sf.TRAFFIC_HARMFUL} x {100 * (1 - miss):.1f}% = {sf.TRAFFIC_HARMFUL * (1 - miss):.1f}, rounded up to {caught_t}. The other {missed_t} are missed",
        f"{harmless_t:,} x {100 * alarm:.2f}% = {harmless_t * alarm:.2f}, rounded up to {false_t}",
        f"{caught_t} + {false_t} = {blocked_t}. The harmful share is {caught_t} / {blocked_t} = {100 * caught_t / blocked_t:.1f}%, which rounds to {sf.share(caught_t, blocked_t)}%",
        # the sweep: the text, the alt text of its figure and the rows beside them
        f"at 0.01 no harmful message is missed but {sweep[1][2]} of {harmless} harmless ones are blocked. At 0.99 no harmless message is blocked but "
        f"{sweep[99][1]} of {harmful} harmful ones are missed",
        f"is 0 until about 0.{max(p for p, c in sweep.items() if c[1] == 0):02d} and then climbs to {100 * sweep[99][1] / harmful:.0f}% at 0.99",
        f"starts at {100 * sweep[1][2] / harmless:.1f}% at 0.01 and falls to 0 at about 0.{min(p for p, c in sweep.items() if c[2] == 0):02d}. They cross near 0.{cross:02d}",
        f"At 0.50 the miss rate is {100 * miss:.1f}% and the false-alarm rate is {100 * alarm:.2f}%",
        f"Caught is {harmful} minus missed, and passed is {harmless} minus alarms",
        f"is {min(c[1] + c[2] for c in sweep.values())}, reached at {len(fewest)} thresholds from 0.{min(fewest):02d} to 0.{max(fewest):02d}",
        f"at 0.{cross:02d}, where each is {100 * sf.rates(sweep[cross])[0]:.1f}% ({sweep[cross][1]} of {harmful} and {sweep[cross][2]} of {harmless})",
        f"0.25 misses none of the {harmful} harmful test messages and blocks {sweep[25][2]} of {harmless} harmless ones",
        f"at least half from 0.{half:02d}, where the test messages have no false alarm left, and by then about {sf.at_traffic(sweep[half])[1]} of the {sf.TRAFFIC_HARMFUL} harmful messages are missed",
        "the parts add to the sum only within 0.002",
        f"The {false_alarm} false alarm is {false_alarm} of only {harmless} harmless test messages, so one more or one fewer would move the {sf.TRAFFIC:,}-message count by about "
        f"{harmless_t / harmless:.0f}, and 0 false alarms in {harmless} is not a rate of 0",
        # the rates and counts at 0.50 as the text, the figures' alt text and the questions give them
        f"At 0.50 the toy misses {missed} of {harmful} harmful test messages ({100 * miss:.1f}%) and raises {false_alarm} false alarm among {harmless} harmless ones ({100 * alarm:.2f}%)",
        f"If only 1 message in {sf.TRAFFIC // sf.TRAFFIC_HARMFUL} is harmful",
        f"at threshold 0.50 on {len(test)} test messages: {missed} missed, {caught} caught, {passed} passed correctly, {false_alarm} false alarm",
        f"leaves {missed} harmful messages on the passed side and {false_alarm} harmless message on the blocked side",
        f"Of {caught + false_alarm} blocked test messages, {caught} are harmful ({sf.share(caught, caught + false_alarm)}%). Of about {blocked_t} blocked in {sf.TRAFFIC:,} messages where "
        f"1 in {sf.TRAFFIC // sf.TRAFFIC_HARMFUL} is harmful, about {caught_t} are harmful ({sf.share(caught_t, blocked_t)}%) and about {false_t} are harmless ({sf.share(false_t, blocked_t)}%)",
        f"is {r['losses'][0]:.4f}, which is ln 2", f"Each of the {len(r['losses'])} steps of gradient descent",
        f"every threshold from 0.01 to 0.99. This table shows some of them on the {len(test)} test messages",
        f"for {sf.TRAFFIC:,} messages of which {sf.TRAFFIC_HARMFUL} are harmful", f"to {sf.TRAFFIC:,} messages of which {sf.TRAFFIC_HARMFUL} are harmful. By hand at 0.50:",
        f"Treat the {false_t} with care", f"{len(test)} messages, with {harmless} harmless", f"the {sf.share(caught_t, blocked_t)}% above",
        f"Here 0.50 sits inside the range, 0.{min(fewest):02d} to 0.{max(fewest):02d}, where the test set has its fewest mistakes ({min(c[1] + c[2] for c in sweep.values())})",
        f"many unseen wordings on purpose ({unseen} of {len(test)})",
        f"At 1 harmful message in {sf.TRAFFIC // sf.TRAFFIC_HARMFUL}, the toy filter's false-alarm rate is a low {100 * alarm:.2f}%",
        f"Why are all {len(wrong)} mistakes at 0.50 in wordings the toy never saw",
        f"Also, the {min(c[1] + c[2] for c in sweep.values())} counts both kinds equally on a test set that is 1 harmful in {len(test) // sum(labels)}, and at 1 in {sf.TRAFFIC // sf.TRAFFIC_HARMFUL}",
        f"The {sf.TRAFFIC:,}-message count at 0.50 is {false_t} false alarms",
        f"It comes from {false_alarm} false alarm among only {harmless} harmless test messages. One more or one fewer would move it by about {harmless_t / harmless:.0f}, "
        f"and 0 in {harmless} would not be a rate of 0",
        # the examples, the weight table and the worked tables
        f'"{test[miss_i][0]}", a request for a phone number in slang, scores {scores[miss_i]:.3f} and passes at 0.50. "{test[alarm_i][0]}" scores {scores[alarm_i]:.3f} and is blocked',
        f"All {len(wrong)} mistakes at 0.50 are in wordings", f"None of the {SAME} test messages in wordings it did see is a mistake",
        f'"phone" and "address" each appear in {in_train("phone", 1)} harmful training messages and in no harmless one. "the" appears in {in_train("the", 0)} of the '
        f"{sum(1 - y for _, y in train)} harmless training messages and in none of the harmful ones",
        *(f"| {word} | {weight(word):.2f} |" for word in ("555", "live", "phone", "address", "street", "home", "me", "number", "is", "great", "the")),
        *worked,
        f'"phone" carries +{weight("phone"):.3f}', f"push the letter-pair total to +{sum(v for n, v in sf.explain(test[alarm_i][0], r)[0] if n.startswith('c:')):.3f}",
        f'has "address" at +{weight("address"):.2f}, but "the" at {weight("the"):.2f} and its other features bring the sum to {z_p:.3f} and the score to {scores[pass_i]:.3f}',
        # the Check-yourself answers
        f"There are {harmless_t // sf.TRAFFIC_HARMFUL} harmless messages for every harmful one. 1.25% of {harmless_t:,} harmless messages is about {false_t}, more than the "
        f"roughly {caught_t} harmful messages caught, so about {false_t} of the roughly {blocked_t} blocked messages ({sf.share(false_t, blocked_t)}%) are harmless",
        f"the fewest mistakes is {min(c[1] + c[2] for c in sweep.values())}", f"{false_alarm} false alarm among only {harmless} harmless test messages",
    ]
    missing = [text for text in quoted if text not in lesson]
    assert not missing, missing


def test_the_figures_are_current_and_follow_the_house_rules(r):
    """Guards a figure left stale after the filter changed, a colour the site rules out, and the two hand-drawn figures
    (flow.svg, outcomes.svg), whose numbers must be the counts and the score the code gives."""
    for name, svg in make_figures.figures().items():
        assert (LESSON / "figs" / name).read_text() == svg, f"{name} is out of date: run make_figures.py"
    for path in (LESSON / "figs").glob("*.svg"):
        svg = path.read_text()
        assert not re.search(r"0072B2|56B4E9|D55E00|<script|<image|href=|https?://(?!www\.w3\.org)", svg, flags=re.I), path.name
        assert "<title>" in svg and "<desc>" in svg and 'fill="#ffffff"' in svg
    caught, missed, false_alarm, passed = r["sweep"][50]
    caught_t, _, false_t, _ = sf.at_traffic(r["sweep"][50])                        # traffic counts are expected counts, so the figure says "about"
    rare = (LESSON / "figs" / "rare.svg").read_text()
    assert all(f"about {n}" in rare for n in (caught_t + false_t, caught_t, false_t)) and f"about {caught_t + false_t} of every" in rare
    outcomes = (LESSON / "figs" / "outcomes.svg").read_text()
    assert [int(n) for n in re.findall(r">(\d+)</text>", outcomes)] == [missed, caught, passed, false_alarm]
    assert f"{len(r['labels'])} test messages, threshold 0.50" in outcomes
    X, _ = sf.matrix([("hey i live at 36 oak street", 1)], r["vocab"])
    score = sf.scores(X, r["w"], r["b"])[0]
    flow = (LESSON / "figs" / "flow.svg").read_text()
    assert score >= 0.5 and f"Here: {score:.2f}" in flow and "Here: 0.50" in flow