import argparse
from collections import Counter
import math
import random
import re

TRAIN = [
    "the library opens at nine",
    "the library closes at five",
    "the reading room opens at nine",
    "the reading room closes at five",
    "the media desk opens at ten",
    "the media desk closes at four",
    "ask the media desk for a scanner",
    "ask the front desk for a card",
    "return the book at the front desk",
]
TEST = [
    "the library opens at ten",
    "the reading room closes at four",
    "ask the media desk for a card",
    "return the magazine at the front desk",
]
START, STOP, UNKNOWN = "<start>", "<end>", "<unk>"


def tokenize(text):
    return re.findall(r"[a-z]+", text.casefold())


def train(documents):
    unigram = Counter()
    bigram = {}
    for text in documents:
        tokens = [START, *tokenize(text), STOP]
        for current, following in zip(tokens, tokens[1:]):
            unigram[following] += 1
            bigram.setdefault(current, Counter())[following] += 1
    vocabulary = sorted({*unigram, UNKNOWN})
    return unigram, bigram, vocabulary


def distribution(counts, vocabulary, alpha):
    denominator = counts.total() + alpha * len(vocabulary)
    return {token: (counts[token] + alpha) / denominator for token in vocabulary}


def perplexity(documents, rows, vocabulary):
    loss = count = 0
    for text in documents:
        tokens = [START, *(token if token in vocabulary else UNKNOWN for token in tokenize(text)), STOP]
        for current, following in zip(tokens, tokens[1:]):
            loss -= math.log(rows[current][following])
            count += 1
    return math.exp(loss / count)


def sample(rows, vocabulary, seed):
    rng = random.Random(seed)
    current = START
    words = []
    for _ in range(18):
        current = rng.choices(vocabulary, weights=list(rows[current].values()))[0]
        if current == STOP:
            break
        words.append(current)
    return " ".join(words) or "(end at first step)"


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train a count-based next-token model on fictional notes.")
    parser.add_argument("--alpha", type=float, default=0.25, help="Smoothing count added to each option")
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if not math.isfinite(args.alpha) or not 0 < args.alpha <= 1_000:
        parser.error("Choose a positive smoothing count up to 1000.")
    unigram, bigram, vocabulary = train(TRAIN)
    baseline = distribution(unigram, vocabulary, args.alpha)
    contexts = [START, *(token for token in vocabulary if token != STOP)]
    unigram_rows = {token: baseline for token in contexts}
    bigram_rows = {token: distribution(bigram.get(token, Counter()), vocabulary, args.alpha)
                   for token in contexts}
    probabilities = bigram_rows["at"]
    print(f"Training vocabulary: {len(vocabulary)} possible next tokens")
    print("After 'at':", ", ".join(f"{token} {score:.3f}" for token, score in
                                 sorted(probabilities.items(), key=lambda item: -item[1])[:5]))
    print(f"Unigram test perplexity: {perplexity(TEST, unigram_rows, vocabulary):.2f}")
    print(f"Bigram test perplexity: {perplexity(TEST, bigram_rows, vocabulary):.2f}")
    print("Bigram sample:", sample(bigram_rows, vocabulary, args.seed))