"""Sampling on top of lesson 4's toy: temperature, top-p and a seeded draw, using nothing but the standard library.

Lesson 4's toy ends each pass with a list of possible next tokens and how often each followed the text it looked at.
This file turns such a list into chances, reshapes them with a temperature, cuts the unlikely tail with top-p, and
draws one token with a seeded generator. The draw is a step of its own, after the model has given its chances.
widget.html runs the same rules in JavaScript; test_widget.py checks that the two agree.

The order is temperature first, then top-p. A provider may order them differently.
"""
import importlib.util
import sys
from pathlib import Path

LESSON_4 = Path(__file__).resolve().parents[2] / "why-answers-appear-word-by-word" / "code"
SLACK = 1e-9           # rounding slack when a running total is compared with p, so Python and JavaScript cut alike
WARMUP = 4             # random numbers thrown away at the start: a small seed gives a tiny first number
DRAWS = 20             # how many picks one press of Draw 20 times makes


def _load_toy():
    """Lesson 4's toy_gen.py, loaded from its own folder so this lesson uses the very same corpus, tables and generator."""
    spec = importlib.util.spec_from_file_location("lesson4_toy_gen", LESSON_4 / "toy_gen.py")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


toy = _load_toy()

# The rows the widget offers. A context is the run of earlier tokens whose followers the toy counted: three tokens
# for the two passes lesson 4 shows, one token for the other two.
ROWS = [
    {"id": "wait", "name": "After “Wait”", "context": ["Wait"]},
    {"id": "the", "name": "After “the”", "context": ["the"]},
    {"id": "pass1", "name": "Tea, pass 1: “make tea ?”", "context": ["make", "tea", "?"]},
    {"id": "pass11", "name": "Tea, pass 11: “bag in a”", "context": ["bag", "in", "a"]},
]


def row_of(model, context):
    """The followers of a run of 1 to 3 tokens, as the toy counted them: (ranked, total). Most to least common, a tie
    to the token that sorts first, which is the order `chances` in lesson 4 uses."""
    followers = model[len(context)][tuple(context)]
    return sorted(followers.items(), key=lambda item: (-item[1], item[0])), sum(followers.values())


def as_chances(ranked, total):
    """[(token, count)] and a total to [(token, chance)]: a token's chance is its count out of the total."""
    return [(token, count / total) for token, count in ranked]


def with_temperature(chances, temperature):
    """Raise each chance to the power 1 / temperature, then scale the chances back to a total of 1. The same rule as
    `with_temperature` in lesson 12. Below 1 it favours the top token, above 1 it evens the tokens out."""
    if temperature <= 0:
        raise ValueError("temperature must be above 0; use the top token for 'no randomness'")
    powered = [(token, chance ** (1 / temperature)) for token, chance in chances]
    total = 0.0
    for _, chance in powered:                       # one number at a time, in order, as widget.html adds them
        total += chance
    return [(token, chance / total) for token, chance in powered]


def top_p(chances, p):
    """Keep the fewest top tokens whose chances add up to at least p, then scale them back to a total of 1.
    `chances` must be ranked from the likeliest down. Tokens with equal chances keep the order they came in.
    A p of 1 keeps every token, even one whose chance is too small to show."""
    if p >= 1:
        return list(chances)
    kept, running = [], 0.0
    for token, chance in chances:
        kept.append((token, chance))
        running += chance
        if running >= p - SLACK:
            break
    return [(token, chance / running) for token, chance in kept]


def reshape(chances, temperature, p):
    """What the draw sees: temperature first, then top-p. Tokens top-p cuts are missing from the result."""
    return top_p(with_temperature(chances, temperature), p)


def choose(chances, u):
    """The token whose slice of 0 to 1 holds u, for u from 0 up to but below 1. Slices follow the order given."""
    running = 0.0
    for token, chance in chances:
        running += chance
        if u < running:
            return token
    return chances[-1][0]                           # u at or past the rounded-off total


def start_state(seed):
    """The generator state a seed starts from: lesson 4's `seed_state`, then WARMUP numbers thrown away.
    Without that, a small seed makes a first number near 0 and the first draw always takes the top token."""
    state = toy.seed_state(seed)
    for _ in range(WARMUP):
        state, _ = toy.seed_step(state)
    return state


def draw(chances, state, n=DRAWS):
    """n draws from the same chances. Returns (the picks, the generator state after them), so the next press goes on."""
    picks = []
    for _ in range(n):
        state, u = toy.seed_step(state)
        picks.append(choose(chances, u))
    return picks, state


def tally(picks, tokens):
    """How often each token in `tokens` came up in `picks`, in the order of `tokens`."""
    return [picks.count(token) for token in tokens]


def generate(model, prompt, state, temperature=1.0, p=1.0, cap=toy.CAP):
    """Write a reply with lesson 4's loop, but each pass draws from the reshaped chances. Returns (the tokens added,
    end mark included, and the generator state after them), so pressing regenerate goes on from where the last
    reply stopped. The same state and settings give the same reply."""
    tokens = toy.tokenize(prompt)
    reply = []
    for _ in range(cap):
        _, ranked, total = toy.chances(model, tokens)
        if not ranked:
            break
        state, u = toy.seed_step(state)
        token = choose(reshape(as_chances(ranked, total), temperature, p), u)
        reply.append(token)
        if token == toy.END:
            break
        tokens.append(token)
    return reply, state


def reply_text(reply):
    """The reply as a reader sees it (see `reply_of` in lesson 4, which takes whole passes)."""
    return toy.reply_of([{"token": token} for token in reply])


def first_split(a, b):
    """The position of the first token where two token lists differ, or None when one is the start of the other."""
    for i, (x, y) in enumerate(zip(a, b)):
        if x != y:
            return i
    return None if len(a) == len(b) else min(len(a), len(b))


def replies(model, prompt, count, seed, temperature=1.0, p=1.0, cap=toy.CAP):
    """`count` replies written one after another from one generator, as if the reader pressed regenerate that often."""
    state, out = start_state(seed), []
    for _ in range(count):
        reply, state = generate(model, prompt, state, temperature, p, cap)
        out.append(reply)
    return out


def corpus_replies(lines, prompt):
    """The reply part of every corpus line that starts with `prompt`, as text (a set, so a repeated line counts once)."""
    return {line[len(prompt):].strip() for line in lines if line.startswith(prompt)}