"""A small safety filter you can read end to end, and the arithmetic of its threshold.

The toy filter blocks one kind of message: asking for, or sharing, personal contact details (a home address,
a phone number, where someone lives) in a school or club chat. Every message is made up here from short
templates, with a fixed seed, so every number in the lesson can be reproduced. A real filter covers many
kinds of harm and is trained on far more text, but it works the same way: a score per message, compared
with a threshold.

    1. make the messages (train and test), some test wordings never seen in training
    2. turn each message into features: words, and pairs of letters inside each word
    3. train a logistic regression with plain gradient descent (NumPy)
    4. score the test messages (0 to 1), then block those at or above the threshold
    5. count the four outcomes, sweep the threshold, and repeat the counts for traffic where harm is rare

Python 3 standard library and NumPy. The widget (../widget.html) repeats the counting in JavaScript;
test_widget.py runs both and compares them.
"""
from __future__ import annotations

import re

import numpy as np

SEED = 15
OPENING_PCT = 50            # the threshold the widget opens on, in hundredths
TRAFFIC = 10_000            # the second traffic mix: 10,000 messages ...
TRAFFIC_HARMFUL = 100       # ... of which 100 are harmful (1 in 100)

# --- the made-up messages ------------------------------------------------------------------------
# {w} slots are filled from the lists below. Harmless messages share words with harmful ones on purpose.
LISTS = {
    "thing": ["home address", "address", "phone number", "cell number", "home phone"],
    "street": ["oak", "maple", "pine", "elm", "cedar", "lake", "hill"],
    "num": [str(n) for n in range(2, 60)],
    "digits": [f"555 01{n:02d}" for n in range(100)],
    "day": ["monday", "tuesday", "wednesday", "thursday", "friday", "saturday"],
    "event": ["practice", "the meeting", "the quiz", "the bake sale", "the trip", "the game", "choir"],
    "room": ["gym", "library", "cafeteria", "art room", "music room", "front office", "science lab", "auditorium"],
    "item": ["jacket", "charger", "notebook", "water bottle", "hoodie", "calculator"],
    "stuff": ["notes", "schedule", "photos", "slides", "worksheet", "results", "link"],
    "food": ["pizza", "snacks", "cupcakes", "juice", "fruit"],
    "who": ["she", "he", "your friend"],
    "poss": ["her", "his", "your friend's"],
    "pre": ["", "", "", "hey ", "hi ", "ok so ", "quick question "],
}
HARMFUL = [
    "what's your {thing}?", "can i get your {thing}?", "send me your {thing}", "where do you live?",
    "which street do you live on?", "my phone number is {digits}", "i live at {num} {street} street",
    "call me on {digits}", "text me at {digits}", "my home address is {num} {street} street",
    "where does {who} live?", "what's {poss} {thing}?",
]
HARMLESS = [
    "the home team plays on {day}", "my number on the team is {num}",
    "can someone send me the {stuff}?", "what time is {event}?", "where is the {room}?", "i forgot my {item}",
    "see you at {event}", "thanks for the {stuff}", "who is bringing the {food}?", "what is the homework for {day}?",
    "can i get the {stuff}?", "i will send the {stuff} after {event}", "your {item} looks great",
    "what's your favorite {food}?", "where do you want to eat?", "the game is live on {day}",
    "you can reach the {room} at the desk", "i will bring the {food} to {event}", "{event} is moving to {day}",
    "does anyone have a spare {item}?", "great job at {event}", "i can help with the {stuff} on {day}",
]
# Wordings used only in the test set: slang, misspellings, and harmless messages that reuse harmful words.
UNSEEN_HARMFUL = [
    "{do} ur digits", "whats ur {adr}", "where do u stay", "which street r u on", "{do} me ur number",
    "what area do u stay in?", "ur place is on which street?", "{do} me ur digits", "what street do u stay on",
]
UNSEEN_HARMLESS = [
    "what's the {room}'s address?", "the school phone number is on the website", "send me the address for {event}",
    "can i get the {room} phone number?", "what's the address of the {room}?", "the home game address is on the poster",
    "can i get the address for {event}?", "what's the {room}'s phone number?", "my phone battery died", "my phone is in my locker",
]
LISTS.update({"do": ["drop", "give", "send"], "adr": ["adress", "addy", "number"]})


class Rng:
    """xorshift32: the same numbers on every Python version, which random.choice does not promise."""

    def __init__(self, seed: int):
        self.state = seed

    def pick(self, items):
        s = self.state
        s ^= (s << 13) & 0xFFFFFFFF
        s ^= s >> 17
        s ^= (s << 5) & 0xFFFFFFFF
        self.state = s
        return items[s % len(items)]


def fill(template: str, rng: Rng) -> str:
    return re.sub(r"\{(\w+)\}", lambda m: rng.pick(LISTS[m.group(1)]), template)


def draw(rng: Rng, harmful_templates, harmless_templates, n_harmful, n_harmless, seen):
    """n messages of each kind: a list of (text, 1 if harmful else 0). No message repeats, and none is in `seen`."""
    out = []
    for templates, n, label in ((harmful_templates, n_harmful, 1), (harmless_templates, n_harmless, 0)):
        got = tries = 0
        while got < n:
            tries += 1
            if tries > 100 * n:
                raise ValueError("the templates make too few different messages")
            core = fill(rng.pick(templates), rng)
            if core not in seen:
                seen.add(core)
                out.append((rng.pick(LISTS["pre"]) + core, label))
                got += 1
    return out


def make_data():
    """(train, test): train is 100 harmful + 200 harmless; test is 40 + 80, of which 16 + 36 are in wordings the filter never saw."""
    rng, seen = Rng(SEED), set()
    train = draw(rng, HARMFUL, HARMLESS, 100, 200, seen)
    same = draw(rng, HARMFUL, HARMLESS, 24, 44, seen)
    unseen = draw(rng, UNSEEN_HARMFUL, UNSEEN_HARMLESS, 16, 36, seen)
    return train, same + unseen


# --- features ------------------------------------------------------------------------------------

def features(text: str) -> set[str]:
    """Each word ("w:address") and each pair of letters inside a word, with ^ and $ marking its ends ("c:ad")."""
    feats = set()
    for word in re.findall(r"[a-z0-9]+", text.lower()):
        feats.add("w:" + word)
        padded = "^" + word + "$"
        feats.update("c:" + padded[i:i + 2] for i in range(len(padded) - 1))
    return feats


def vocabulary(train):
    """Features seen in training, in a fixed order. A feature never seen in training gets no weight."""
    return {name: i for i, name in enumerate(sorted(set().union(*(features(t) for t, _ in train))))}


def matrix(messages, vocab):
    X = np.zeros((len(messages), len(vocab)))
    for row, (text, _) in enumerate(messages):
        for name in features(text):
            if name in vocab:
                X[row, vocab[name]] = 1.0
    return X, np.array([label for _, label in messages], dtype=float)


# --- training and scoring ------------------------------------------------------------------------

def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-z))


def train_filter(X, y, steps=300, rate=0.5, l2=0.003):
    """Gradient descent on the average log-loss plus a small penalty on the weights. Returns weights, bias, loss per step."""
    w, b, losses = np.zeros(X.shape[1]), 0.0, []
    for _ in range(steps):
        p = sigmoid(X @ w + b)
        losses.append(float(-np.mean(y * np.log(p) + (1 - y) * np.log(1 - p))))
        err = p - y
        w -= rate * (X.T @ err / len(y) + l2 * w)
        b -= rate * err.mean()
    return w, b, losses


def scores(X, w, b):
    """One score per message, 0 to 1, rounded to three decimals (what the widget shows)."""
    return [round(float(p), 3) for p in sigmoid(X @ w + b)]


# --- counting ------------------------------------------------------------------------------------

def counts(labels, score_list, pct):
    """Block a message when its score is at least pct/100. Returns (caught, missed, falsely blocked, passed)."""
    t = pct / 100
    caught = sum(1 for y, s in zip(labels, score_list) if y == 1 and s >= t)
    blocked_harmless = sum(1 for y, s in zip(labels, score_list) if y == 0 and s >= t)
    harmful = sum(labels)
    return caught, harmful - caught, blocked_harmless, len(labels) - harmful - blocked_harmless


def rates(c):
    """(miss rate, false-alarm rate, precision): missed / harmful, falsely blocked / harmless, caught / blocked."""
    caught, missed, false_alarm, passed = c
    blocked = caught + false_alarm
    return missed / (caught + missed), false_alarm / (false_alarm + passed), (caught / blocked if blocked else None)


def sweep(labels, score_list):
    """Counts at every threshold from 0.01 to 0.99."""
    return {pct: counts(labels, score_list, pct) for pct in range(1, 100)}


def crossing(table):
    """The first threshold where the miss rate is at least the false-alarm rate: where the two curves cross."""
    return next(pct for pct in range(1, 100) if rates(table[pct])[0] >= rates(table[pct])[1])


def at_traffic(c, total=TRAFFIC, harmful=TRAFFIC_HARMFUL):
    """The same miss and false-alarm rates applied to `total` messages of which `harmful` are harmful.

    Whole numbers, rounding halves up, so Python and the widget agree exactly.
    """
    caught, missed, false_alarm, passed = c
    n_harmful, n_harmless = caught + missed, false_alarm + passed
    caught_t = (caught * harmful + n_harmful // 2) // n_harmful
    false_t = (false_alarm * (total - harmful) + n_harmless // 2) // n_harmless
    return caught_t, harmful - caught_t, false_t, total - harmful - false_t


# --- everything together ---------------------------------------------------------------------------

def build():
    """Train the filter and score the test messages. Returns a dict the tests, the numbers and the widget share."""
    train, test = make_data()
    vocab = vocabulary(train)
    X, y = matrix(train, vocab)
    w, b, losses = train_filter(X, y)
    Xt, yt = matrix(test, vocab)
    test_scores = scores(Xt, w, b)
    return {"train": train, "test": test, "vocab": vocab, "w": w, "b": b, "losses": losses,
            "labels": [int(v) for v in yt], "scores": test_scores, "sweep": sweep([int(v) for v in yt], test_scores)}


# --- the words the widget shows ----------------------------------------------------------------------

def find(test, core: str) -> int:
    """Index of the test message that ends with `core` (a random greeting may come before it)."""
    return next(i for i, (text, _) in enumerate(test) if text.endswith(core))


MISSED_EXAMPLE = "drop ur digits"                    # harmful, in a wording the training messages never had
FALSE_ALARM_EXAMPLE = "my phone battery died"       # harmless, but "phone" is a word harmful messages use
PASSED_EXAMPLE = "what's the library's address?"    # harmless and passed, although it uses the word "address"


def noun(n: int) -> str:
    return "message" if n == 1 else "messages"


def share(part: int, whole: int) -> int:
    """part / whole as a whole percent, halves rounded up."""
    return (200 * part + whole) // (2 * whole)


def view(labels, score_list, pct: int, traffic: bool = False) -> dict:
    """What the widget shows at one threshold: the four counts, a sentence, and the share of blocked messages that were harmful.

    traffic=False counts the test messages; traffic=True applies the same rates to TRAFFIC messages of which TRAFFIC_HARMFUL are harmful,
    so its counts are expected values and its sentences say "about".
    The widget offers thresholds 1 to 99; at every one of them the filter blocks at least 14 test messages, so `blocked` is never 0.
    """
    c = counts(labels, score_list, pct)
    if traffic:
        c = at_traffic(c)
    caught, missed, false_alarm, passed = c
    t = f"0.{pct:02d}"
    if traffic:
        says = (f"Out of {TRAFFIC:,} messages, {TRAFFIC_HARMFUL} harmful: at {t} the filter would miss about {missed} "
                f"and block about {false_alarm:,} harmless {noun(false_alarm)}.")
    else:
        says = f"At {t} the filter misses {missed} harmful {noun(missed)} and blocks {false_alarm} harmless {noun(false_alarm)}."
    blocked = caught + false_alarm
    if traffic:
        precision = f"Of the {blocked:,} or so messages it would block, about {caught:,} are harmful ({share(caught, blocked)}%)."
    else:
        precision = f"Of the {blocked:,} messages it blocks, {caught:,} are harmful ({share(caught, blocked)}%)."
    return {"counts": list(c), "says": says, "precision": precision}


def readout(text: str, label: int, score: float, pct: int) -> str:
    """The one line shown for a chosen message."""
    blocked = score >= pct / 100
    verdict = {(1, True): "caught", (1, False): "a miss", (0, True): "a false alarm", (0, False): "correct"}[(label, blocked)]
    return f"“{text}” ({'harmful' if label else 'harmless'}) scores {score:.3f}, so it is {'blocked' if blocked else 'passed'}: {verdict}."


def explain(text: str, result: dict):
    """The features of one message the filter has weights for, those weights, the features it never saw, and the sum before the sigmoid."""
    vocab, w = result["vocab"], result["w"]
    feats = features(text)
    known = sorted((name, float(w[vocab[name]])) for name in feats if name in vocab)
    unseen = sorted(name for name in feats if name not in vocab)
    return known, unseen, result["b"] + sum(v for _, v in known)
