"""A toy language model that only counts which words follow which, using nothing but the standard library.

widget.html runs the same arithmetic in JavaScript; test_widget.py checks that the two agree.

The model reads the last two words of a prompt and asks three questions in turn, from most specific to
least: what followed these two words in the corpus, what followed the last word, and (when nothing at all
has been seen) any word of the corpus at random. The three answers are blended with the Witten-Bell
rule: the more different words have followed a context, the less that context is trusted.
"""
import re
from pathlib import Path

CORPUS = Path(__file__).resolve().parent / "corpus.txt"
START = "<s>"
TOKEN = re.compile(r"[A-Za-z0-9']+|[.,!?]")      # ASCII only, so Python and JavaScript split alike
IDK = "I don't know"


def load_corpus():
    return CORPUS.read_text(encoding="utf-8")


def tokenize(text):
    """Words and the marks . , ! ? as separate tokens. Capital letters count: 'Peru' is not 'peru'."""
    return TOKEN.findall(text)


def lines_of(text):
    """The non-empty lines of a corpus. Each line is one sentence."""
    return [line for line in text.split("\n") if line.strip()]


def train(lines):
    """Count what follows each pair of words and each single word. Every line starts after two START marks."""
    pairs, singles, vocab = {}, {}, set()
    for line in lines:
        words = [START, START] + tokenize(line)
        vocab.update(words[2:])
        for i in range(2, len(words)):
            for table, context in ((pairs, tuple(words[i - 2:i])), (singles, (words[i - 1],))):
                followers = table.setdefault(context, {})
                followers[words[i]] = followers.get(words[i], 0) + 1
    return {"pairs": pairs, "singles": singles, "vocab": sorted(vocab)}


def next_word_probs(model, prompt):
    """Probability of every word of the corpus coming next after the prompt. They add up to 1."""
    words = [START, START] + tokenize(prompt)
    vocab = model["vocab"]
    if not vocab:
        return {}
    probs = {w: 1 / len(vocab) for w in vocab}
    for table, context in ((model["singles"], tuple(words[-1:])), (model["pairs"], tuple(words[-2:]))):
        followers = table.get(context)
        if not followers:
            continue
        seen = sum(followers.values())
        trust = seen / (seen + len(followers))      # Witten-Bell: seen / (seen + kinds of follower)
        probs = {w: trust * followers.get(w, 0) / seen + (1 - trust) * probs[w] for w in vocab}
    return probs


def with_temperature(probs, temperature):
    """Raise every probability to 1/temperature and scale back to total 1. Below 1 it sharpens, above 1 it flattens."""
    powered = {w: p ** (1 / temperature) for w, p in probs.items()}
    total = sum(powered.values())
    return {w: p / total for w, p in powered.items()}


def ranked(probs):
    """(word, probability) from most to least likely. A tie goes to the word that sorts first."""
    return sorted(probs.items(), key=lambda item: (-item[1], item[0]))


def answer(probs, threshold=0.0):
    """The most likely word, or IDK when even that word has a probability below the threshold."""
    if not probs:
        return IDK
    word, p = ranked(probs)[0]
    return word if p >= threshold else IDK


def model_with_extra(lines, true_sentence, k):
    """The corpus with k more copies of one true sentence."""
    return train(lines + [true_sentence] * k)


def smallest_k(lines, prompt, true_sentence, truth, threshold=0.0, temperature=1.0, limit=50):
    """The fewest added copies of the true sentence after which the model answers `truth`, being at least
    `threshold` sure of it. None when `limit` copies are not enough."""
    for k in range(limit + 1):
        probs = with_temperature(next_word_probs(model_with_extra(lines, true_sentence, k), prompt), temperature)
        if answer(probs, threshold) == truth:
            return k
    return None