"""A toy that writes a reply one token at a time, using nothing but the standard library.

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

The toy reads the text so far (the question and the reply written up to now), looks at its last three tokens,
and asks what followed those three tokens in its corpus. If it never saw them it looks at the last two, then the
last one, then at all tokens together. The share of each follower is its chance. One pass of the loop gives one
chance list and one token, which is added to the text before the next pass.
"""
import re
from math import floor
from pathlib import Path

CORPUS = Path(__file__).resolve().parent / "corpus.txt"
START = "<s>"
END = "<end>"
TOKEN = re.compile(r"[A-Za-z0-9']+|[.,!?]")      # ASCII only, so Python and JavaScript split alike
ORDER = 3                                         # how many earlier tokens one pass looks at
CAP = 40                                          # the longest reply the widget allows
SEED_MODULUS = 2147483647                         # 2**31 - 1, for the seeded pick (Park-Miller)
SEED_FACTOR = 48271


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


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


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


def train(lines):
    """Count what followed every run of 0 to ORDER tokens. Each line is padded with START marks and ends with END."""
    tables = [{} for _ in range(ORDER + 1)]
    for line in lines:
        tokens = [START] * ORDER + tokenize(line) + [END]
        for i in range(ORDER, len(tokens)):
            for k in range(ORDER + 1):
                followers = tables[k].setdefault(tuple(tokens[i - k:i]), {})
                followers[tokens[i]] = followers.get(tokens[i], 0) + 1
    return tables


def chances(model, tokens):
    """What one pass produces: the run of tokens it looked at and every possible next token with its count.

    Returns (looked_at, ranked, total). `ranked` is [(token, count)] from most to least likely; a tie goes to the
    token that sorts first. The chance of a token is its count divided by the total."""
    text = [START] * ORDER + list(tokens)
    for k in range(ORDER, -1, -1):
        looked_at = tuple(text[len(text) - k:])
        followers = model[k].get(looked_at)
        if followers:
            ranked = sorted(followers.items(), key=lambda item: (-item[1], item[0]))
            return looked_at, ranked, sum(followers.values())
    return (), [], 0                                # an empty corpus


def seed_state(seed):
    """The first state of the seeded pick: a whole number from 1 to 2**31 - 2."""
    return seed % (SEED_MODULUS - 1) + 1


def seed_step(state):
    """One move of the seeded generator (Park-Miller). Returns (next state, a number from 0 up to but below 1)."""
    state = state * SEED_FACTOR % SEED_MODULUS
    return state, (state - 1) / (SEED_MODULUS - 1)


def pick(ranked, total, u=None):
    """The token the pass adds. With u=None the likeliest; with u in [0, 1) a draw in which each token gets
    its count out of the total. Only whole numbers are compared, so Python and JavaScript agree."""
    if u is None:
        return ranked[0][0]
    r = floor(u * total)
    for token, count in ranked:
        if r < count:
            return token
        r -= count
    return ranked[-1][0]


def generate(model, prompt, cap=30, seed=None):
    """Write a reply one pass at a time. Returns one dict per pass: the text before it, what it looked at, its
    chances, and the token it added. Stops after the END mark or when the reply reaches `cap` tokens.
    With seed=None every pass takes the likeliest token; with a seed it draws by chance."""
    tokens = tokenize(prompt)
    state = seed_state(seed) if seed is not None else None
    passes = []
    for _ in range(cap):
        looked_at, ranked, total = chances(model, tokens)
        if not ranked:
            break
        u = None
        if state is not None:
            state, u = seed_step(state)
        token = pick(ranked, total, u)
        passes.append({"before": list(tokens), "looked_at": list(looked_at), "ranked": ranked, "total": total, "token": token})
        if token == END:
            break
        tokens.append(token)
    return passes


def reply_of(passes):
    """The reply as the reader sees it: the tokens without the END mark, with marks . , ! ? attached to the word before."""
    out = ""
    for p in passes:
        if p["token"] == END:
            break
        out += p["token"] if p["token"] in ".,!?" else (" " if out else "") + p["token"]
    return out


def positions_read(prompt_tokens, new_tokens, cached):
    """How many token positions the network runs over to write `new_tokens` tokens after a prompt of `prompt_tokens`.

    Without a cache every pass re-reads the whole text so far. With the cache of earlier work (the attention cache)
    the first pass reads the prompt and every later pass reads only the newest token."""
    if new_tokens <= 0:
        return 0
    if cached:
        return prompt_tokens + (new_tokens - 1)
    return sum(prompt_tokens + i for i in range(new_tokens))